fix a minor bug in test sync_pipeline_task_gen_test.py

PiperOrigin-RevId: 620958204
tensorflow · Apr 1, 2024 · 8787340 · 8787340
1 parent 6844762
commit 8787340
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tfx/orchestration/experimental/core/sync_pipeline_task_gen_test.py b/tfx/orchestration/experimental/core/sync_pipeline_task_gen_test.py
@@ -904,6 +904,7 @@ def test_restart_node_cancelled_due_to_stopping_with_foreach(self):
         num_new_executions=1,
         num_active_executions=1,
         ignore_update_node_state_tasks=True,
+        pipeline=pipeline,
     )
     self.assertEqual(stats_gen_node_uid, stats_gen_task.node_uid)
     self.assertIsInstance(stats_gen_task, task_lib.ExecNodeTask)
@@ -987,6 +988,7 @@ def test_restart_node_cancelled_due_to_fail_with_foreach(self):
         num_new_executions=2,
         num_active_executions=2,
         ignore_update_node_state_tasks=True,
+        pipeline=pipeline,
     )
     self.assertEqual(stats_gen_node_uid, stats_gen_task.node_uid)
     self.assertIsInstance(stats_gen_task, task_lib.ExecNodeTask)