Hey all, Has anybody had success with mocking `flow.with_options`? I have an async flow `refresh_da...

Avinash Santhanagopalan

03/14/2024, 3:57 PM

Hey all, Has anybody had success with mocking

flow.with_options

? I have an async flow

refresh_data_for_plugin

that I used to patch like this

Copy code

patch("refresh_data_for_plugin", new_callable=AsyncMock)

But after adding

refresh_data_for_plugin.with_options

I am not able to patch it in a similar way

Copy code

patch("refresh_data_for_plugin.with_options", new_callable=AsyncMock)

I’m getting

TypeError: 'coroutine' object is not callable

Any help appreciated!

Dominic Tarro

03/14/2024, 5:21 PM

Task.with_options

will return an awaitable. It is not an awaitable itself.

Dominic Tarro

03/14/2024, 5:22 PM

Instead, set a

return_value

for

with_options

that returns your

AsyncMock

👀 1

Avinash Santhanagopalan

03/14/2024, 5:37 PM

Thanks @Dominic Tarro but I tried that and got. I think without patching

refresh_data_for_plugin

first I wouldn’t be able to patch

with_options

Copy code

TypeError: An asyncio.Future, a coroutine or an awaitable is required

But I’m not too familiar with async pytesting so not sure how to proceed

Dominic Tarro

03/14/2024, 5:43 PM

May have to adjust the details, but I remember doing something like this

Copy code

patch("refresh_data_for_plugin", Mock(with_options=Mock(return_value=...)))

Let me know if that works

👀 1

Dominic Tarro

03/14/2024, 5:56 PM

Copy code

from unittest.mock import Mock, patch

from prefect import task


@task
def unmocked_func(*args, **kwargs):
    print("I'm not mocked!")


def test_unmocked_func():
    def mocked_func():
        print("I'm mocked!")

    with patch(
        "test.unmocked_func",
        Mock(
            with_options=Mock(
                return_value=mocked_func
            )
        )
    ) as mocked:
        unmocked_func.with_options()()

thank you 1

Avinash Santhanagopalan

03/14/2024, 6:29 PM

Thanks that worked! A small follow up question, how can I

assert_has_calls

unmocked_func

Dominic Tarro

03/14/2024, 6:31 PM

Ah, instead of returning an actual function you can return that AsyncMock.

Copy code

async_mock = AsyncMock()
with patch(..., ...return_value=async_mock...):
    # stuff
    async_mock.assert_has_calls()

33 Views

Open in Slack

Previous Next

Prefect Community

Bring your towel and join one of the fastest growing data communities. Welcome to our second-generation open source orchestration platform, a completely rethought approach to dataflow automation.