https://prefect.io logo
Join Slack
Powered by
<@ULVA73B9P> I'm writing a unit test using pytest...
# ask-marvin
c
@Marvin I'm writing a unit test using pytest a 
flow
that calls a 
task
using Prefect's 
.map()
over a collection of objects. How can I mock the task execution?
m
To mock the execution of a task while using Prefect's 
.map()
in your unit tests, you can use the original task function by accessing it with 
.fn
. Prefect offers facilities to aid in testing flows and tasks, including context managers and fixtures for integrating with testing frameworks such as pytest. Here's a concise approach to achieve this: 1. Access the Original Task Function: By using 
.fn
, you can call the underlying function of a task directly, making it easy to mock or test in isolation.
Copy code
python
   from prefect import flow, task
   from unittest.mock import patch

   @task
   def my_task(x):
       return x * 2

   @flow
   def my_flow(values):
       results = my_task.map(values)
       return results

   def test_my_flow():
       with patch('my_task.fn', return_value=42):
           result = my_flow([1, 2, 3])
           assert result == [42, 42, 42]
2. Use a Prefect Test Harness: Prefect provides utilities like the 
prefect_test_harness
, which allows running flows and tasks against a temporary local database. This can be used as a pytest fixture:
Copy code
python
   import pytest
   from prefect import flow
   from prefect.testing.utilities import prefect_test_harness

   @pytest.fixture(autouse=True, scope="session")
   def prefect_test_fixture():
       with prefect_test_harness():
           yield

   def test_my_flow_with_fixture():
       # Here, you can run your flow and verify behavior under test conditions.
For more detailed information on testing workflows and tasks within Prefect, you can explore the Prefect 2.x Testing Guide and the 3.x Testing Guide. Remember, even in the robotic world, a task tested is a flow half-completed!
c
What if it uses a Prefect Future? Can I mock a future.result?
@Marvin How to unit test a task that uses prefect future
m
When unit testing tasks in Prefect that use 
PrefectFuture
, you'll want to carefully manage your mocks to simulate task completion and access the underlying results properly. Here's a concise way to approach this: 1. Using 
.result()
Method: When you obtain a 
PrefectFuture
by submitting a task, you can mock the behavior of the 
.result()
method to return a desired value. This is useful for tasks waiting on upstream task results. 2. Mocking with 
unittest.mock
: You can use 
unittest.mock
to patch methods or objects, including futures, to simulate their behavior during tests. Here's a brief example:
Copy code
python
from unittest.mock import Mock, patch
from prefect import task, flow

@task
def my_task(x):
    return x * 2

@flow
def my_flow():
    future = my_task.submit(10)
    result = future.result()
    return result

def test_my_flow():
    with patch('path.to.my_task.submit') as mock_submit:
        mock_future = Mock()
        mock_future.result.return_value = 20
        mock_submit.return_value = mock_future

        result = my_flow()
        assert result == 20
In this example: - We use 
patch
to mock the 
submit
method of the 
my_task
to return a mock future. - The mock future's 
result()
method is then set to return the expected result of the computation. - The flow 
my_flow()
is executed, and you assert the result matches the mocked outcome. For further insights, you can refer to the following resources: - Prefect 2.x Concepts on Task Runners - Prefect 3.x Task Runner Documentation Hopefully, 
PrefectFutures
will behave today as predictably as the sun rises—unless you're on a planet with erratic rotations, of course!
n
hi @Charlotte! feel free to share your test here and explain what the unit test is meant to show lots of times in our tests we just actually run the tasks against the 
prefect_test_harness
fixture like marvin mentioned, but if you really need to mock out map you could
20 Views
Open in Slack
PreviousNext
Powered by