test workflow endpoint doesn't seem to always send it back as complete but we can poll for that

Author: chrishagglund-ship-it

tests/integration/client/orkes/test_orkes_clients.py

@@ -26,7 +26,7 @@
 from conductor.client.workflow.conductor_workflow import ConductorWorkflow
 from conductor.client.workflow.executor.workflow_executor import WorkflowExecutor
 from conductor.client.workflow.task.simple_task import SimpleTask
-from tests.integration.retry_helpers import retry_scenario
+from tests.integration.retry_helpers import retry_scenario, TERMINAL_WORKFLOW_STATES
 
 SUFFIX = str(uuid())
 WORKFLOW_NAME = 'IntegrationTestOrkesClientsWf_' + SUFFIX
@@ -394,24 +394,28 @@ def __test_unit_test_workflow(self):
         execution = self.workflow_client.test_workflow(testRequest)
         assert execution != None
 
-        print(
-            f"[test_workflow] workflow_id={getattr(execution, 'workflow_id', None)} "
-            f"status={execution.status} "
-            f"task_count={len(execution.tasks or [])}"
-        )
-        for t in (execution.tasks or []):
-            print(
-                f"[test_workflow]   task ref={t.reference_task_name} "
-                f"type={t.task_type} status={t.status} "
-                f"retried={getattr(t, 'retry_count', None)}"
-            )
+        # There appears to be no guarantee about it actually coming back
+        # complete, because it happens (often) that it does not. So we accept a
+        # COMPLETED result immediately, but if it isn't COMPLETED yet we don't
+        # fail: we poll the workflow and wait for it to reach a terminal state.
         if execution.status != "COMPLETED":
+            #need to fix the print statement below to have the wf id and the status
             print(
-                f"[test_workflow] non-terminal output={getattr(execution, 'output', None)}"
+                f"[test_workflow] workflow_id={getattr(execution, 'workflow_id', None)} status={execution.status} (was expecting COMPLETED - but will poll for that now)"
             )
+            workflow_id = getattr(execution, "workflow_id", None)
+            polled = (
+                self.__poll_workflow_until_complete(workflow_id)
+                if workflow_id else None
+            )
+            if polled is not None:
+                execution = polled
 
         # Ensure workflow is completed successfully
-        assert execution.status == "COMPLETED"
+        assert execution.status == "COMPLETED", (
+            f"workflow expected to be COMPLETED, but received {execution.status}, "
+            f"workflow_id: {getattr(execution, 'workflow_id', None)}"
+        )
 
         # Ensure the inputs were captured correctly
         assert execution.input["loanAmount"] == testRequest.input["loanAmount"]
@@ -463,6 +467,35 @@ def __test_unit_test_workflow(self):
         # Workflow output takes the latest iteration output of a loopOver task.
         assert execution.output["phoneNumberValid"]
 
+    def __poll_workflow_until_complete(self, workflow_id, timeout_seconds=60,
+                                       poll_interval=2):
+        """Poll ``workflow_id`` until it reaches a terminal state or the timeout
+        passes, returning the last observed Workflow (or None if it could never
+        be fetched). Transient poll errors are logged and retried within the
+        timeout so a slow-but-eventually-complete run isn't reported as a bare
+        failure.
+        """
+        deadline = time.monotonic() + timeout_seconds
+        workflow = None
+        while True:
+            try:
+                workflow = self.workflow_client.get_workflow(
+                    workflow_id, include_tasks=True)
+            except Exception as e:  # transient blip against the shared server
+                print(f"[test_workflow] error polling {workflow_id}: {e}")
+            status = getattr(workflow, "status", None)
+            if status in TERMINAL_WORKFLOW_STATES:
+                print(f"[test_workflow] {workflow_id} reached {status}")
+                return workflow
+            if time.monotonic() >= deadline:
+                print(
+                    f"[test_workflow] {workflow_id} still {status} after "
+                    f"{timeout_seconds}s; giving up wait"
+                )
+                return workflow
+            print(f"[test_workflow] {workflow_id} still {status}; waiting")
+            time.sleep(poll_interval)
+
     def __test_unregister_workflow_definition(self):
         self.metadata_client.unregister_workflow_def(WORKFLOW_NAME, 1)