https://prefect.io logo
Title
l

Luke Orland

12/06/2019, 4:51 PM
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.
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:
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

Chris White

12/06/2019, 5:17 PM
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

Luke Orland

12/06/2019, 5:19 PM
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

Chris White

12/06/2019, 5:30 PM
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

Luke Orland

12/06/2019, 5:43 PM
ok, great, this is working:
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

Chris White

12/06/2019, 5:56 PM
Awesome!!
@Marvin archive “How to mock a task run when testing a flow”