example StepDecorator test

Author: ryan-williams

metaflow/tests/flows/step_counter_flow.py

@@ -0,0 +1,19 @@
+from metaflow.plugins import STEP_DECORATORS
+
+from metaflow.tests.flows import NewLinearFlow
+from metaflow.tests.step_counter import StepCounter
+
+STEP_DECORATORS.append(StepCounter)
+
+
+class StepCounterFlow(NewLinearFlow):
+    """Example flow that can receive `--with step_counter` from the CLI
+
+    In order for `--with my_decorator` to work on the CLI, `my_decorator` needs to added to `STEP_DECORATORS`, as above.
+    I don't know of a good way to accomplish that other than making a trivial flow (`StepCounterFlow` here) that just
+    mixes in the target flow (`NewLinearFlow`) in a file that appends the decorator (`StepCounter`) to `STEP_DECORATORS.
+
+    TODO: formalize an API for registering step- and flow-decorators without having to define a new flow like this.
+    """
+
+    pass

metaflow/tests/step_counter.py

@@ -0,0 +1,13 @@
+from metaflow.decorators import StepDecorator
+
+
+class StepCounter(StepDecorator):
+    name = "step_counter"
+
+    def task_post_step(
+        self, step_name, flow, graph, retry_count, max_user_code_retries
+    ):
+        if not hasattr(flow, "step_count"):
+            flow.step_count = 0
+        flow.step_count += 1
+        print("step count: %d" % flow.step_count)

metaflow/tests/test_decorators.py

@@ -0,0 +1,20 @@
+from metaflow.tests.flows.step_counter_flow import StepCounterFlow
+from metaflow.tests.utils import run
+
+
+def test_step_count():
+    flow = StepCounterFlow
+    data = run(
+        flow,
+        args=(
+            "run",
+            "--with",
+            "step_counter",
+        ),
+    )
+    assert data == {
+        "a": 111,
+        "b": 222,
+        "checked": True,
+        "step_count": 5,
+    }