[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?

• Yes

Have you read the Code of Conduct?

• Yes

[welcome]

Thank you for opening your first issue here! 🛠

[runllm]

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!