test: add regression tests for metaflow-mage QA findings

Adds tests that catch bugs found in the deep-QA run of metaflow-mage
(run 20260315-193000). Tests are written to fail against the unfixed
code and pass after fixes are applied.

test_dag.py:
- test_foreach, test_multibody_foreach: assert len(tasks) == N to catch
  silent foreach_count=1 fallback (D-FOREACH-1)
- test_nested_foreach_2x2: 2-outer x 2-inner structure asserting all 4
  combinations appear in outer_join and 4 inner tasks exist; the minimal
  1-inner case in test_nested_foreach passes even with D-NESTED-1 bug

test_compliance.py:
- test_timeout_minutes_enforced: verifies @timeout(minutes=1) actually kills
  the step (D-TIMEOUT-1); skipped on remote backends same as test_timeout_enforcement
- test_run_param_not_dropped: deploys flow with retry_count param, triggers
  with retry_count=42, asserts value arrives at task runtime

test_utils.py:
- wait_for_deployed_run: assert triggered_run.run is not None, catching
  run_id mismatches between deployer and init block in every deployer test

New flows:
- flows/basic/timeout_minutes_flow.py: @timeout(minutes=1) on 120s sleep
- flows/basic/reserved_param_flow.py: retry_count Parameter for param passthrough test
- flows/dag/nested_foreach_2x2_flow.py: 2x2 nested foreach (groups x items)

Author: Nissan Pow

test/ux/core/flows/basic/reserved_param_flow.py

@@ -0,0 +1,31 @@
+"""Flow with a parameter named 'retry_count' to test parameter passing correctness.
+
+When deployed to Mage, parameters were silently dropped when trigger variables
+had None values or when JSON serialization lost the value. This flow verifies
+that parameter values arrive correctly at task runtime.
+"""
+
+from metaflow import FlowSpec, Parameter, step, project
+
+
+@project(name="reserved_param_flow")
+class ReservedParamFlow(FlowSpec):
+    retry_count = Parameter(
+        "retry_count",
+        default=0,
+        type=int,
+        help="Number of retries (named retry_count to test parameter passing)",
+    )
+
+    @step
+    def start(self):
+        self.stored_retry_count = self.retry_count
+        self.next(self.end)
+
+    @step
+    def end(self):
+        pass
+
+
+if __name__ == "__main__":
+    ReservedParamFlow()

test/ux/core/flows/basic/timeout_minutes_flow.py

@@ -0,0 +1,34 @@
+"""Flow with @timeout(minutes=1) on a step that sleeps 120 seconds.
+
+Used to verify that @timeout(minutes=N) is correctly parsed and enforced —
+not silently treated as @timeout(seconds=0).
+"""
+
+import time
+
+from metaflow import FlowSpec, step, timeout, project
+
+
+@project(name="timeout_minutes_flow")
+class TimeoutMinutesFlow(FlowSpec):
+    """Step with @timeout(minutes=1) sleeps 2 minutes — must be killed after 1 minute."""
+
+    @step
+    def start(self):
+        self.next(self.slow)
+
+    @timeout(minutes=1)
+    @step
+    def slow(self):
+        # Sleep 2 minutes — should be killed after 1 minute timeout.
+        time.sleep(120)
+        self.done = True
+        self.next(self.end)
+
+    @step
+    def end(self):
+        pass
+
+
+if __name__ == "__main__":
+    TimeoutMinutesFlow()

test/ux/core/flows/dag/nested_foreach_2x2_flow.py

@@ -0,0 +1,43 @@
+"""Nested foreach flow with 2 outer x 2 inner items — tests full 2x2 fanout."""
+
+import os
+
+from metaflow import FlowSpec, step, project
+
+
+@project(name="nested_foreach_2x2_flow")
+class NestedForeach2x2Flow(FlowSpec):
+    @step
+    def start(self):
+        self.execution_env = os.environ.get("KUBERNETES_SERVICE_HOST", "")
+        self.groups = ["x", "y"]
+        self.next(self.outer, foreach="groups")
+
+    @step
+    def outer(self):
+        self.group = self.input
+        self.items = [1, 2]
+        self.next(self.inner, foreach="items")
+
+    @step
+    def inner(self):
+        self.result = "%s-%d" % (self.group, self.input)
+        self.next(self.inner_join)
+
+    @step
+    def inner_join(self, inputs):
+        self.inner_results = sorted([i.result for i in inputs])
+        self.next(self.outer_join)
+
+    @step
+    def outer_join(self, inputs):
+        self.all_results = sorted([r for i in inputs for r in i.inner_results])
+        self.next(self.end)
+
+    @step
+    def end(self):
+        pass
+
+
+if __name__ == "__main__":
+    NestedForeach2x2Flow()

test/ux/core/test_compliance.py

@@ -359,3 +359,100 @@ def test_timeout_enforcement(exec_mode, decospecs, compute_env, tag, scheduler_c
         "Run should have failed because the 'slow' step exceeds its "
         "@timeout(seconds=5), but it succeeded. Timeout enforcement may be broken."
     )
+
+
+# ---------------------------------------------------------------------------
+# test_timeout_minutes_enforced
+#
+# WHY: _get_timeout_seconds only read 'seconds' attribute, silently ignoring
+# 'minutes'.  @timeout(minutes=1) produced no timeout at all — the step ran
+# indefinitely.  This test verifies that minute-based timeouts are actually
+# enforced (D-TIMEOUT-1).
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.compliance
+@pytest.mark.scheduler_only
+@pytest.mark.skip(
+    reason="@timeout enforcement on remote backends (argo/sfn/airflow) is not "
+    "reliable — the run may hang instead of failing. Needs backend-specific "
+    "timeout mechanisms (e.g. activeDeadlineSeconds for k8s). See #XXXX."
+)
+def test_timeout_minutes_enforced(
+    exec_mode, decospecs, compute_env, tag, scheduler_config
+):
+    """WHY: _get_timeout_seconds only read 'seconds' attribute, silently ignoring 'minutes'.
+    @timeout(minutes=1) produced no timeout at all.
+    This test verifies that minute-based timeouts are actually enforced (D-TIMEOUT-1).
+    """
+    if exec_mode != "deployer":
+        pytest.skip("compliance test requires deployer mode")
+
+    test_unique_tag = f"test_compliance_timeout_minutes_{exec_mode}"
+    combined_tags = tag + [test_unique_tag]
+
+    tl_args = {
+        "env": compute_env,
+        "decospecs": decospecs,
+    }
+
+    deployed_flow = deploy_flow_to_scheduler(
+        flow_name="basic/timeout_minutes_flow.py",
+        tl_args=tl_args,
+        scheduler_args={"cluster": scheduler_config.cluster},
+        deploy_args={"tags": combined_tags, **(scheduler_config.deploy_args or {})},
+        scheduler_type=scheduler_config.scheduler_type,
+    )
+
+    run = wait_for_deployed_run_allow_failure(deployed_flow)
+
+    assert not run.successful, (
+        "@timeout(minutes=1) was NOT enforced — the step ran for 2+ minutes without being killed. "
+        "Check that _get_timeout_seconds correctly computes minutes*60+seconds."
+    )
+
+
+# ---------------------------------------------------------------------------
+# test_run_param_not_dropped
+#
+# WHY: Parameters were silently dropped when trigger variables dict had None
+# values or when JSON serialization lost the value.  Verify parameter values
+# arrive correctly at task runtime.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.compliance
+@pytest.mark.scheduler_only
+def test_run_param_not_dropped(
+    exec_mode, decospecs, compute_env, tag, scheduler_config
+):
+    """WHY: Parameters were silently dropped when trigger variables dict had None values
+    or when JSON serialization lost the value. Verify parameter values arrive correctly.
+    """
+    if exec_mode != "deployer":
+        pytest.skip("compliance test requires deployer mode")
+
+    test_unique_tag = f"test_compliance_run_param_not_dropped_{exec_mode}"
+    combined_tags = tag + [test_unique_tag]
+
+    tl_args = {
+        "env": compute_env,
+        "decospecs": decospecs,
+    }
+
+    deployed_flow = deploy_flow_to_scheduler(
+        flow_name="basic/reserved_param_flow.py",
+        tl_args=tl_args,
+        scheduler_args={"cluster": scheduler_config.cluster},
+        deploy_args={"tags": combined_tags, **(scheduler_config.deploy_args or {})},
+        scheduler_type=scheduler_config.scheduler_type,
+    )
+
+    run = wait_for_deployed_run(deployed_flow, run_kwargs={"retry_count": 42})
+
+    assert run.successful, "Run was not successful"
+    assert run["start"].task.data.stored_retry_count == 42, (
+        "Expected retry_count=42, got %r. "
+        "Parameter may have been dropped or not passed correctly."
+        % run["start"].task.data.stored_retry_count
+    )

test/ux/core/test_dag.py

@@ -36,6 +36,12 @@ def test_foreach(exec_mode, decospecs, compute_env, tag, scheduler_config):
     )
 
     assert run.successful, "Run was not successful"
+    # Verify exact fanout count — catches silent foreach_count=1 fallback (D-FOREACH-1)
+    process_tasks = list(run["process"].tasks())
+    assert len(process_tasks) == 3, (
+        "Expected 3 foreach tasks for items=[1,2,3], got %d. "
+        "This may indicate foreach_count fell back to 1." % len(process_tasks)
+    )
     assert run["join"].task.data.results == [
         2,
         4,
@@ -56,6 +62,10 @@ def test_multibody_foreach(exec_mode, decospecs, compute_env, tag, scheduler_con
     )
 
     assert run.successful, "Run was not successful"
+    process_tasks = list(run["process"].tasks())
+    assert len(process_tasks) == 3, "Expected 3 foreach process tasks, got %d" % len(
+        process_tasks
+    )
     assert run["join"].task.data.results == [
         3,
         5,
@@ -153,3 +163,46 @@ def test_nested_foreach(exec_mode, decospecs, compute_env, tag, scheduler_config
         "x-1",
         "y-1",
     ], "Nested foreach all_results didn't match"
+
+
+def test_nested_foreach_2x2(exec_mode, decospecs, compute_env, tag, scheduler_config):
+    """Verify nested foreach with 2 outer x 2 inner items — catches D-NESTED-1 semantic bug."""
+    from metaflow.exception import MetaflowException
+
+    try:
+        run = execute_test_flow(
+            flow_name="dag/nested_foreach_2x2_flow.py",
+            exec_mode=exec_mode,
+            decospecs=decospecs,
+            tag=tag,
+            scheduler_config=scheduler_config,
+            test_name="nested_foreach_2x2",
+            tl_args_extra={"env": compute_env},
+        )
+    except (MetaflowException, Exception) as e:
+        msg = str(e).lower()
+        if exec_mode == "deployer" and (
+            "not supported" in msg or "not yet supported" in msg
+        ):
+            pytest.skip(
+                f"{scheduler_config.scheduler_type} does not support nested foreach: {e}"
+            )
+        raise
+
+    assert run.successful, "Run was not successful"
+    # Must have all 4 combinations: x-1, x-2, y-1, y-2
+    assert run["outer_join"].task.data.all_results == [
+        "x-1",
+        "x-2",
+        "y-1",
+        "y-2",
+    ], (
+        "Expected 4 results from 2x2 nested foreach, got: %s. "
+        "This may indicate nested_foreach_join is not aggregating all outer items correctly."
+        % run["outer_join"].task.data.all_results
+    )
+    # Verify inner task count: 4 inner tasks total (2 outer x 2 inner)
+    inner_tasks = list(run["inner"].tasks())
+    assert (
+        len(inner_tasks) == 4
+    ), "Expected 4 inner tasks for 2x2 foreach, got %d" % len(inner_tasks)

test/ux/core/test_utils.py

@@ -182,7 +182,12 @@ def wait_for_deployed_run(
         print("Waiting for run to start...")
         time.sleep(polling_interval)
 
-    print(f"Run {triggered_run.run.id} started")
+    run = triggered_run.run
+    assert run is not None, (
+        "triggered_run.run returned None — run_id mismatch between deployer and init block. "
+        "Check that pipeline_run_id kwarg is injected by the scheduler."
+    )
+    print(f"Run {run.id} started")
 
     while not triggered_run.run.finished:
         if time.time() - start_time > timeout: