l

    Luke Orland

    2 years ago
    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.
    Chris White

    Chris White

    2 years ago
    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

    2 years ago
    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...
    Chris White

    Chris White

    2 years ago
    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

    2 years ago
    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!
    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=)
    )
    Chris White

    Chris White

    2 years ago
    Awesome!!
    @Marvin archive “How to mock a task run when testing a flow”
    Marvin

    Marvin

    2 years ago