wait for wf more robustly in test

Author: chrishagglund-ship-it

tests/integration/test_comprehensive_e2e.py

@@ -195,36 +195,66 @@ def test_02_start_workers(self):
         self.__class__.workers_started = True
         self.assertTrue(self.workers_started)
 
-    def test_03_execute_workflow(self):
-        """Execute workflow and verify completion."""
-        print("\n" + "="*80 + "\nTEST 3: Execute Workflow\n" + "="*80)
-        
-        self.assertTrue(self.workers_started, "Workers must be started first")
-        
-        workflow_client = OrkesWorkflowClient(self.config)
+    EXPECTED_TASK_COUNT = 5
+
+    def _run_workflow_to_terminal(self, workflow_client, timeout_s=90):
+        """Start the e2e workflow and wait until it is genuinely terminal with
+        all expected tasks present. Returns (wf_id, workflow_or_None). A None
+        workflow (or a non-terminal / short task list) signals a timed-out run
+        the caller can retry, rather than asserting against a half-scheduled
+        workflow (which is what produced the flaky "4 != 5 tasks" failure).
+        """
         req = StartWorkflowRequest()
         req.name = 'e2e_comprehensive_test'
         req.version = 1
         req.input = {}
-        
+
         wf_id = workflow_client.start_workflow(start_workflow_request=req)
         print(f"✓ Started workflow: {wf_id}")
-        
-        # Wait for completion
-        for i in range(30):
+
+        deadline = time.time() + timeout_s
+        wf = None
+        while time.time() < deadline:
             wf = workflow_client.get_workflow(wf_id, include_tasks=True)
-            print(f"  Status: {wf.status} - ({i*2}s)")
-            if wf.status in ['COMPLETED', 'FAILED']:
-                break
-            for task in wf.tasks:
+            print(f"  Status: {wf.status} - tasks={len(wf.tasks or [])}")
+            if wf.status in ('COMPLETED', 'FAILED') \
+                    and len(wf.tasks or []) == self.EXPECTED_TASK_COUNT:
+                return wf_id, wf
+            for task in (wf.tasks or []):
                 print(f'task {task.task_def_name} : {task.status}')
             time.sleep(1)
+
+        return wf_id, wf
+
+    def test_03_execute_workflow(self):
+        """Execute workflow and verify completion."""
+        print("\n" + "="*80 + "\nTEST 3: Execute Workflow\n" + "="*80)
 
-        final_wf = workflow_client.get_workflow(wf_id, include_tasks=True)
+        self.assertTrue(self.workers_started, "Workers must be started first")
 
+        workflow_client = OrkesWorkflowClient(self.config)
+
+        # Each attempt starts a fresh workflow, so retrying a pathologically
+        # slow run (cold workers / slow CI) is safe and self-contained. This
+        # keeps a one-off timeout from forcing a manual CI job re-run.
+        final_wf = None
+        for attempt in range(3):
+            wf_id, wf = self._run_workflow_to_terminal(workflow_client, timeout_s=90)
+            if wf is not None and wf.status in ('COMPLETED', 'FAILED') \
+                    and len(wf.tasks or []) == self.EXPECTED_TASK_COUNT:
+                final_wf = wf
+                break
+            print(f"attempt {attempt + 1}: wf={wf_id} "
+                  f"status={getattr(wf, 'status', None)} "
+                  f"tasks={len(getattr(wf, 'tasks', []) or [])}; retrying")
+
         # Assertions
-        self.assertIsNotNone(final_wf)
-        self.assertEqual(len(final_wf.tasks), 5, "Should have 5 tasks")
+        self.assertIsNotNone(
+            final_wf,
+            "workflow never reached a terminal state with all "
+            f"{self.EXPECTED_TASK_COUNT} tasks after retries")
+        self.assertEqual(len(final_wf.tasks), self.EXPECTED_TASK_COUNT,
+                         f"Should have {self.EXPECTED_TASK_COUNT} tasks")
 
         # Verify each task
         tasks_by_ref = {t.reference_task_name: t for t in final_wf.tasks}