https://prefect.io
Join Slack
I am developing a workflow, and I'm having trouble writing tests that mock out the `run()` method of...
# ask-community
l
I am developing a workflow, and I'm having trouble writing tests that mock out the 
run()
method of tasks in the workflow.
I've tried 
unittest.mock.patch
, 
unittest.mock.patch.object
, also followed the examples in the prefect codebase with `monkeypatch. With the following test (run with 
pytest
), I'm finding that the query is being executed for real and not patched.
Copy code
from unittest.mock import MagicMock

from myworkflow import flow


def test_flow(monkeypatch):

    MockSnowflakeQuery = MagicMock()
    mock_run = MagicMock()
    MockSnowflakeQuery.run = mock_run
    monkeypatch.setattr(
        "myworkflow.SnowflakeQuery", MockSnowflakeQuery,
    )

    flow.run()

    args, kwargs = mock_run.call_args
    assert args == ()
    assert kwargs['query'].startswith('COPY INTO')
The flow has a task:
Copy code
mytask = SnowflakeQuery(**prefect.context.secrets.SNOWFLAKE).run(
        query=myquery
    )
My high-level goal is to be able to develop the flow and its tasks while mocking out all I/O to the various databases and services that will be involved.
c
Hey @Luke Orland ! Because the Flow has already been created, and the task has already been instantiated, your mock isn’t having any effect - try mocking “my_task” and see if that solves it
l
my_task, in this case, is an instance of SnowflakeQuery, correct?
oh gotcha
I'm not seeing any difference.
Perhaps patching isn't the best way to develop workflows...
c
Actually, in the code snippet above it looks like you already called the run method of the task?
You should be patching the instance of the task that is associated to the flow, and it should work as you expect
l
ok, great, this is working:
Copy code
from unittest.mock import patch

from kepler.recommendations.workflow import flow, mytask


@patch.object(mytask, 'run')
def test_flow(mock_mytask_run):
    flow.run()

    mock_mytask_run.assert_called()
Thanks!
💯 1
I removed the 
.run(...)
call from the SnowflakeQuery() instance in the assignment of 
mytask
(and move the 
query=
kwarg from 
SnowflakeQuery().run(query=...)
into 
SnowflakeQuery(query=)
)
c
Awesome!!
@Marvin archive “How to mock a task run when testing a flow”
m
Open in Slack
PreviousNext
Powered by

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.

Powered by