@Marvin Can you find me any examples of testing flow code which has hooks for on_completion or on_failure? I would like to mock these methods ideally. How can I do this?

To test flow code with hooks such as

on_completion

or

on_failure

, you can mock these methods following some guidance I found in the Prefect documentation and external resources. However, none of the results directly mention how to specifically mock these hooks, but we can derive how to do it using generic Python testing and mocking principles, as well as some insights from the documentation. Here's what we can do: 1. Use the guide for state change hooks to understand how to define hooks: State Change Hooks. 2. Create the hooks that you want to invoke on completion or failure of the flow and then pass them to the

@flow

decorator as shown in the guides, using the parameters

on_completion

and/or

on_failure

. Examples:

python
from prefect import flow

def my_flow_hook(flow, flow_run, state):
    # define what should happen when the flow completes or fails
    pass

@flow(on_completion=[my_flow_hook], on_failure=[my_flow_hook])
def my_workflow():
    # flow logic here
    pass

3. Mock the hook functions during testing using the

unittest.mock

library or a similar mocking library in Python. Here is a generic example of how you could mock a function:

python
from unittest.mock import MagicMock
from mymodule import my_workflow

def test_my_workflow():
    # mock the hook functions
    with unittest.mock.patch('mymodule.my_flow_hook', new=MagicMock()) as mock_hook:
        # call the workflow
        my_workflow()
        
        # verify that the hook has been called if applicable
        mock_hook.assert_called()

4. For more information on mocking and testing, refer to the Python

unittest.mock

documentation or other testing frameworks like

pytest

. Unfortunately, I couldn't find a specific code example that demonstrates precisely mocking hooks in Prefect flows. If you need a more detailed guide tailored to Prefect, you might want to consider raising an issue on the Prefect Github or asking in the Prefect community for insights from other developers who have specifically tackled this problem. Additionally, if you want to learn more about how different states are handled in flows, you can refer to the documentation on Flow Run States. For mocking context-specific variables, including perhaps aspects of flow or task state, see the guide on Runtime Context, which briefly touches on mocking values.

👋 @Blake

forest gump wave

do you want to remove the hooks for the test? or just mock the functions?

Ideally i just want to assert that they each get called. But I would settle for removing them in the test as well

how about

caplog

?

import pytest
from prefect import flow


def success_hook(flow, flow_run, state):
    print("hook called")


def fail_hook(flow, flow_run, state):
    print("hook called")


@flow(on_completion=[success_hook], on_failure=[fail_hook])
def foo(should_fail: bool = False):
    if should_fail:
        raise ValueError("foo failed")


def test_hooks_called(caplog):
    with caplog.at_level("INFO"):
        foo()
        assert "Running hook 'success_hook'" in caplog.text

    with caplog.at_level("INFO"):
        with pytest.raises(ValueError, match="foo failed"):
            foo(should_fail=True)
        assert "Running hook 'fail_hook'" in caplog.text