Add composition tests for single-step flows

Cover the scenarios raised in PR review for @step(start=True, end=True)
single-step flows: Config descriptors, stacked step decorators, and
FlowMutators.

Unit tests (test/unit/test_graph_structure.py):
- Config descriptor is registered via _get_parameters.
- @Retry + @resources stack correctly on the only step.
- FlowMutator registers at class-definition time (pre_mutate only fires
  via the CLI layer, so execution is covered by the integration test).

Integration tests (test/unit/graph_inference/, new flow files + fixtures):
- Config-bearing single-step flow runs end-to-end; config value flows
  through to the end task's artifact.
- Stacked @retry/@resources flow runs; _graph_info records both
  decorators on the only step.
- FlowMutator-decorated flow runs; pre_mutate lands @Retry on the sole
  step.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: talsperre

test/unit/graph_inference/conftest.py

@@ -35,3 +35,23 @@ def flow_fixture(request):
 custom_named_card_run = pytest.fixture(scope="session")(
     create_flow_fixture("CustomNamedCardFlow", "custom_named_card_flow.py")
 )
+
+single_step_with_config_run = pytest.fixture(scope="session")(
+    create_flow_fixture(
+        "SingleStepWithConfigFlow", "single_step_with_config_flow.py"
+    )
+)
+
+single_step_with_stacked_decos_run = pytest.fixture(scope="session")(
+    create_flow_fixture(
+        "SingleStepWithStackedDecosFlow",
+        "single_step_with_stacked_decos_flow.py",
+    )
+)
+
+single_step_with_flow_mutator_run = pytest.fixture(scope="session")(
+    create_flow_fixture(
+        "SingleStepWithFlowMutatorFlow",
+        "single_step_with_flow_mutator_flow.py",
+    )
+)

test/unit/graph_inference/flows/single_step_with_config_flow.py

@@ -0,0 +1,15 @@
+from metaflow import FlowSpec, step, Config
+
+
+class SingleStepWithConfigFlow(FlowSpec):
+    """Single-step flow wired to a Config, verifies Configs resolve end-to-end."""
+
+    cfg = Config("cfg", default_value={"x": 7})
+
+    @step(start=True, end=True)
+    def only(self):
+        self.v = self.cfg["x"]
+
+
+if __name__ == "__main__":
+    SingleStepWithConfigFlow()

test/unit/graph_inference/flows/single_step_with_flow_mutator_flow.py

@@ -0,0 +1,22 @@
+from metaflow import FlowSpec, step, retry, FlowMutator
+
+
+class AddRetryToOnly(FlowMutator):
+    """Adds @retry to every step. Verifies mutators reach the sole step."""
+
+    def pre_mutate(self, mutable_flow):
+        for _, s in mutable_flow.steps:
+            s.add_decorator(retry, deco_kwargs={"times": 1})
+
+
+@AddRetryToOnly
+class SingleStepWithFlowMutatorFlow(FlowSpec):
+    """Single-step flow with a FlowMutator applied at the class level."""
+
+    @step(start=True, end=True)
+    def only(self):
+        self.v = 1
+
+
+if __name__ == "__main__":
+    SingleStepWithFlowMutatorFlow()

test/unit/graph_inference/flows/single_step_with_stacked_decos_flow.py

@@ -0,0 +1,15 @@
+from metaflow import FlowSpec, step, retry, resources
+
+
+class SingleStepWithStackedDecosFlow(FlowSpec):
+    """Single-step flow with @retry and @resources stacked on top of @step."""
+
+    @retry(times=1)
+    @resources(cpu=1, memory=256)
+    @step(start=True, end=True)
+    def only(self):
+        self.v = 42
+
+
+if __name__ == "__main__":
+    SingleStepWithStackedDecosFlow()

test/unit/graph_inference/test_graph_inference.py

@@ -246,3 +246,50 @@ def test_default_card_includes_custom_graph_endpoints(custom_named_card_run):
     assert set(dag_data["steps"]) == {"begin", "middle", "finish"}
     assert "start" not in dag_data["steps"]
     assert "end" not in dag_data["steps"]
+
+
+# ---------------------------------------------------------------------------
+# Composition: single-step flows with Config, stacked decorators, FlowMutator
+# ---------------------------------------------------------------------------
+
+
+def test_single_step_with_config_completes(single_step_with_config_run):
+    """Config-bearing single-step flow runs to completion."""
+    assert single_step_with_config_run.successful
+    assert single_step_with_config_run.finished
+
+
+def test_single_step_with_config_value_flows_to_artifact(single_step_with_config_run):
+    """Config descriptor value is readable from the end task's artifact."""
+    end_task = single_step_with_config_run.end_task
+    assert end_task["v"].data == 7
+
+
+def test_single_step_with_stacked_decos_completes(single_step_with_stacked_decos_run):
+    """Single-step flow with stacked @retry/@resources runs end-to-end."""
+    assert single_step_with_stacked_decos_run.successful
+    assert single_step_with_stacked_decos_run.finished
+
+
+def test_single_step_with_stacked_decos_graph_info(single_step_with_stacked_decos_run):
+    """_graph_info records all stacked decorators on the only step."""
+    graph_info = (
+        single_step_with_stacked_decos_run["_parameters"].task["_graph_info"].data
+    )
+    names = {d["name"] for d in graph_info["steps"]["only"]["decorators"]}
+    assert {"retry", "resources"}.issubset(names)
+
+
+def test_single_step_with_flow_mutator_completes(single_step_with_flow_mutator_run):
+    """FlowMutator-decorated single-step flow runs end-to-end."""
+    assert single_step_with_flow_mutator_run.successful
+    assert single_step_with_flow_mutator_run.finished
+
+
+def test_single_step_with_flow_mutator_applied(single_step_with_flow_mutator_run):
+    """FlowMutator.add_decorator landed @retry on the only step."""
+    graph_info = (
+        single_step_with_flow_mutator_run["_parameters"].task["_graph_info"].data
+    )
+    names = {d["name"] for d in graph_info["steps"]["only"]["decorators"]}
+    assert "retry" in names

test/unit/test_graph_structure.py

@@ -12,7 +12,8 @@
 """
 
 import pytest
-from metaflow import FlowSpec, step, Parameter
+from metaflow import Config, FlowMutator, FlowSpec, step, Parameter, retry, resources
+from metaflow.flowspec import FlowStateItems
 from metaflow.lint import linter, LintWarn
 
 # ---------------------------------------------------------------------------
@@ -133,6 +134,48 @@ def _make_dynamic_single_step_flow():
     )
 
 
+# ---------------------------------------------------------------------------
+# Flow classes: single-step flows composed with configs, decorators, mutators
+# ---------------------------------------------------------------------------
+
+
+class _SingleStepWithConfig(FlowSpec):
+    """Single-step flow with a Config descriptor."""
+
+    cfg = Config("cfg", default_value={"x": 7})
+
+    @step(start=True, end=True)
+    def only(self):
+        self.v = self.cfg["x"]
+
+
+class _SingleStepWithStackedDecos(FlowSpec):
+    """Single-step flow with multiple step decorators stacked."""
+
+    @retry(times=3)
+    @resources(cpu=2, memory=1024)
+    @step(start=True, end=True)
+    def only(self):
+        self.v = 1
+
+
+class _AddRetryMutator(FlowMutator):
+    """Adds @retry to every step. Used to verify mutators reach a single-step flow."""
+
+    def pre_mutate(self, mutable_flow):
+        for _, s in mutable_flow.steps:
+            s.add_decorator(retry, deco_kwargs={"times": 1})
+
+
+@_AddRetryMutator
+class _SingleStepWithFlowMutator(FlowSpec):
+    """Single-step flow with a FlowMutator applied at the class level."""
+
+    @step(start=True, end=True)
+    def only(self):
+        pass
+
+
 # ---------------------------------------------------------------------------
 # Fixtures
 # ---------------------------------------------------------------------------
@@ -487,6 +530,42 @@ def test_backward_compat_name_based():
     assert graph["end"].is_end_step is False
 
 
+# ---------------------------------------------------------------------------
+# Tests: composition with configs, stacked decorators, and flow mutators
+# ---------------------------------------------------------------------------
+
+
+def test_single_step_with_config_descriptor_registered():
+    """Config descriptor is registered on a single-step flow."""
+    graph = _SingleStepWithConfig._graph
+    assert graph.start_step == "only" == graph.end_step
+    names = {name for name, _ in _SingleStepWithConfig._get_parameters()}
+    assert "cfg" in names
+
+
+def test_single_step_with_multiple_step_decorators():
+    """Multiple step decorators stack correctly on a single-step flow."""
+    graph = _SingleStepWithStackedDecos._graph
+    deco_names = {d.name for d in graph["only"].decorators}
+    assert {"retry", "resources"}.issubset(deco_names)
+
+
+def test_single_step_with_flow_mutator_registered():
+    """FlowMutator is registered on a single-step flow at class-definition time.
+
+    pre_mutate only fires when the flow is processed via the CLI layer, so
+    the decorator it adds won't appear on the graph in a unit test. What we
+    can verify here is that the mutator syntax is accepted by a single-step
+    FlowSpec and that it's registered as a flow mutator. End-to-end execution
+    is covered by the matching integration test.
+    """
+    flow_cls = _SingleStepWithFlowMutator._flow_cls
+    graph = flow_cls._graph
+    assert graph.start_step == "only" == graph.end_step
+    mutators = flow_cls._flow_state[FlowStateItems.FLOW_MUTATORS]
+    assert any(isinstance(m, _AddRetryMutator) for m in mutators)
+
+
 # ---------------------------------------------------------------------------
 # Negative-path tests: malformed annotation patterns caught by lint
 # ---------------------------------------------------------------------------