Testing a flow with pytest

You can test Metaflow flows with pytest in two ways: test the logic inside steps, or test the flow end to end.

Testing logic in steps
Testing a flow end to end

A helpful design pattern is to move non-orchestration logic out of your flows and into separate modules, then write unit tests for those functions. For example, say your flow file has logic embedded directly in a step:

my_flow.py

from metaflow import FlowSpec, step

class MyFlow(FlowSpec):

    @step
    def start(self):
        # load the input data
        # transform and validate it
        # write the result
        self.next(self.next_step)

    # rest of flow
    ...

You cannot import and test that logic on its own, because it only exists inside the flow. Move it into a separate module file:

my_module.py

def do_logic():
    # load the input data
    # transform and validate it
    # write the result

Then, update the flow to import the module instead:

my_flow.py

from metaflow import FlowSpec, step

class MyFlow(FlowSpec):

    @step
    def start(self):
        from my_module import do_logic
        do_logic()
        self.next(self.next_step)

    # rest of flow

The logic is identical in both versions of the flow; only its location changes. You can now unit test do_logic independently of the flow, which is the point of the refactor. Separating logic from orchestration makes the code easier to maintain and test, especially when multiple flows or steps share the same logic.

To test that a flow produces the artifact values you expect, use Metaflow’s Runner API from a pytest script. For example, given this flow:

simple_flow.py

from metaflow import FlowSpec, step

class FlowToTest(FlowSpec):

    @step
    def start(self):
        self.x = 0
        self.next(self.end)

    @step
    def end(self):
        self.x += 1

if __name__ == '__main__':
    FlowToTest()

To test that x equals 1 after the flow runs:

Optionally, switch Metaflow profiles

By default, Metaflow creates a profile at ~/.metaflow_config/config.json. To separate test-run data from your actual runs, create a test profile at ~/.metaflowconfig/config_test.json:

{
    "METAFLOW_DEFAULT_DATASTORE": "local"
}

Write the pytest script

Define a test that runs the flow and asserts on the artifact:

test_simple_flow.py

from metaflow import Runner

def test_flow():
    runner = Runner(flow_file="./simple_flow.py", profile="test")
    result = runner.run()
    run_obj = result.run
    assert run_obj.data.x == 1

Run pytest

pytest

The test passes when the artifact value matches:

============================= test session starts ==============================
platform darwin -- Python 3.12.4, pytest-8.2.2, pluggy-1.5.0
plugins: anyio-4.4.0
collected 1 item

test_simple_flow.py .                                                    [100%]

============================== 1 passed in 2.01s ===============================

Anaconda Platform

Getting Started

Platform Guides

Testing a flow with pytest