[Core feature] Support patching subworkflows in unit tests

[charliemoriarty]

Motivation: Why do you think this is important?

Currently, flytekit.testing supports patching tasks in workflows, which allows us to unit test that the sequence of workflow tasks is called as expected. This is helpful for task-only workflows, but for more complex workflows that contain subworkflows, we lack the ability to patch the subworkflows and test that they are similarly called in sequence as expected (without needing to know about the internals of the subworkflows themselves).

These tests would help confirm that workflow sequences fire as expected, and protect against unwanted changes to the execution flow in the codebase.

Goal: What should the final outcome look like, ideally?

Ideally, it would be handy having a test decorator that can be used similarly to the task patching decorator already provided within flytekit.testing:

from flytekit import task, workflow

from flytekit.testing import patch as task_patch
from flytekit.testing import workflow_patch

@task
def final_task():
    return 5

@task
def no_task():
    return 0

@workflow
def my_subworkflow() -> int:
    do_some_complicated_tasks()
    return final_task()

@workflow
def my_workflow(some_condition: bool) -> int:
    result = (
        conditional("my_workflow_switch")
        .if_(some_condition.is_true())
        .then(my_subworkflow())
        .else_()
        .then(no_task())
    )
    return result


@workflow_patch(my_subworkflow)
@task_patch(no_task)
def test_my_workflow_invokes_expected_condition(
    task_mock,
    subworkflow_mock,
):
    my_workflow(True)

    subworkflow_mock.assert_called_once()
    task_mock.assert_not_called()

Describe alternatives you've considered

Without the ability to patch subworkflows, we would need to write unit tests that have knowledge of each workflow in the chain until we reach the task level for patching - this is doable, but breaks the rules of test isolation.

Propose: Link/Inline OR Additional context

No response

Are you sure this issue hasn't been raised already?

• [X] Yes

Have you read the Code of Conduct?

• [X] Yes

[welcome[bot]]

Thank you for opening your first issue here! 🛠

[runllm[bot]]

To support patching subworkflows in unit tests, you can propose this feature on the Flyte GitHub repository. Currently, flytekit.testing does not support direct patching of subworkflows as it does for tasks. This feature would enhance testing capabilities by allowing developers to verify workflow sequences without delving into subworkflow internals.

Would you like more detail on how to contribute this feature to Flyte, including steps to create an issue or pull request on GitHub?

Sources:

• https://discuss.flyte.org/t/13546197/hello-i-have-defined-a-flyte-workflow-which-calls-a-few-flyt/
• https://github.com/flyteorg/flytekit/blob/master/tests/flytekit/unit/remote/test_wrapper_classes.py

---

This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @runllm to ask follow-up questions. See runllm.com for more!

[wayner0628]

Hi @charliemoriarty,

I looked into this issue, and I found that flytekit.testing does support patching subworkflows. I ran your example with the necessary modifications, and it functions well. The only thing to be aware of is that you need to specify the return type of the mocking tasks and ensure that the MagicMock objects return the same type of values.

Code Snippet

from flytekit import task, workflow, conditional
from flytekit.testing import patch

@task
def final_task() -> int:
    return 5

@task
def no_task() -> int:
    return 0

@workflow
def my_subworkflow() -> int:
    # do_some_complicated_tasks()
    return final_task()

@workflow
def my_workflow(some_condition: bool) -> int:
    result = (
        conditional("my_workflow_switch")
        .if_(some_condition.is_true())
        .then(my_subworkflow())
        .else_()
        .then(no_task())
    )
    return result

@patch(my_subworkflow)
@patch(no_task)
def test_my_workflow_invokes_expected_condition(
    task_mock,
    subworkflow_mock,
):
    subworkflow_mock.return_value = 42
    task_mock.return_value = 42
    my_workflow(some_condition=True)

    subworkflow_mock.assert_called_once()
    task_mock.assert_not_called()

if __name__ == "__main__":
    test_my_workflow_invokes_expected_condition()

Output

Does this meet your needs? Feel free to let me know!

[charliemoriarty]

Hi @wayner0628 - apologies for the long delay in reply here. I've done some tests and can confirm that patching the subworkflow as you've suggested above does indeed seem to work! It's a little misleading in my IDE (VSCode) because I get a type error "PythonFunctionWorkflow is incompatible with type PythonFunctionTask" when using the patch decorator or task_mock method - but the tests still pass if I run them 👍 Thank you!