refactored some of the retry stuff to dedupe logic

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, TERMINAL_WORKFLOW_STATES
+from tests.integration.retry_helpers import retry_scenario, wait_for_workflow_terminal
 
 SUFFIX = str(uuid())
 WORKFLOW_NAME = 'IntegrationTestOrkesClientsWf_' + SUFFIX
@@ -471,30 +471,15 @@ 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
+        be fetched). Transient poll errors are swallowed and retried within the
         timeout so a slow-but-eventually-complete run isn't reported as a bare
-        failure.
+        failure. Thin wrapper over the shared ``wait_for_workflow_terminal``.
         """
-        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)
+        return wait_for_workflow_terminal(
+            self.workflow_client, workflow_id,
+            timeout_seconds=timeout_seconds, poll_interval=poll_interval,
+            include_tasks=True, swallow='all',
+            log=lambda msg: print(f"[test_workflow] {msg}"))
 
     def __test_unregister_workflow_definition(self):
         self.metadata_client.unregister_workflow_def(WORKFLOW_NAME, 1)

tests/integration/retry_helpers.py

@@ -140,3 +140,83 @@ def retry_scenario(label, func, *args, deadline=None,
                 deadline - now, transient)
             time.sleep(delay)
             attempt += 1
+
+
+def wait_for_workflow_terminal(
+        client, workflow_id, *,
+        timeout_seconds=None, deadline=None, poll_interval=2,
+        include_tasks=False, terminal_states=TERMINAL_WORKFLOW_STATES,
+        is_terminal=None, swallow='transient',
+        log=None, on_poll=None, on_giveup=None):
+    """Poll ``client.get_workflow(workflow_id, include_tasks=...)`` until the
+    workflow reaches a terminal state (or the wall-clock budget is exhausted),
+    returning the last observed ``Workflow`` (or ``None`` if it could never be
+    fetched). This is the single shared "poll a workflow to terminal" primitive
+    for the integration suites — callers wrap it to preserve their own return
+    shape / side effects.
+
+    Args:
+        client: anything exposing ``get_workflow(workflow_id, include_tasks=...)``
+            (both ``OrkesWorkflowClient`` and ``WorkflowExecutor`` qualify).
+        timeout_seconds: relative budget; an internal ``time.monotonic()``
+            deadline is derived from it. Ignored when ``deadline`` is given.
+        deadline: absolute ``time.monotonic()`` value to stop at. When both this
+            and ``timeout_seconds`` are ``None`` the workflow is polled exactly
+            once (no wait).
+        poll_interval: seconds between polls.
+        include_tasks: passed through to ``get_workflow``.
+        terminal_states: statuses considered terminal (default
+            ``TERMINAL_WORKFLOW_STATES``); used by the default predicate.
+        is_terminal: optional ``predicate(workflow) -> bool`` overriding the
+            default "status in terminal_states" check (e.g. to also gate on an
+            expected task count).
+        swallow: error policy for a failing poll — ``'transient'`` swallows only
+            transient blips (status 0/None) and re-raises real errors;
+            ``'all'`` swallows every exception and keeps polling; ``'none'`` lets
+            any exception propagate.
+        log: optional ``callable(str)`` for progress lines (e.g. ``logger.info``
+            or ``print``). Defaults to ``logger.debug``.
+        on_poll: optional ``callable(workflow)`` invoked after each successful
+            poll, before the terminal check (e.g. to print per-task diagnostics).
+        on_giveup: optional ``callable(workflow)`` invoked once when the budget
+            is exhausted without reaching terminal (e.g. to dump diagnostics).
+    """
+    if log is None:
+        log = logger.debug
+    if is_terminal is None:
+        def is_terminal(wf):
+            return getattr(wf, 'status', None) in terminal_states
+    if deadline is None and timeout_seconds is not None:
+        deadline = time.monotonic() + timeout_seconds
+
+    start = time.monotonic()
+    workflow = None
+    while True:
+        try:
+            workflow = client.get_workflow(
+                workflow_id, include_tasks=include_tasks)
+        except Exception as e:
+            if swallow == 'none':
+                raise
+            if swallow == 'transient' and not is_transient(e):
+                raise
+            log('error polling workflow %s (%.0fs elapsed): %s' % (
+                workflow_id, time.monotonic() - start, e))
+        else:
+            if on_poll is not None:
+                on_poll(workflow)
+            if is_terminal(workflow):
+                log('workflow %s reached %s after %.0fs' % (
+                    workflow_id, getattr(workflow, 'status', None),
+                    time.monotonic() - start))
+                return workflow
+        if deadline is None or time.monotonic() >= deadline:
+            log('workflow %s still %s after %.0fs; giving up wait' % (
+                workflow_id, getattr(workflow, 'status', None),
+                time.monotonic() - start))
+            if on_giveup is not None:
+                on_giveup(workflow)
+            return workflow
+        log('workflow %s still %s; waiting' % (
+            workflow_id, getattr(workflow, 'status', None)))
+        time.sleep(poll_interval)

tests/integration/test_async_lease_extension.py

@@ -38,14 +38,13 @@
 from conductor.client.http.models.task_def import TaskDef
 from conductor.client.http.models.workflow_task import WorkflowTask
 from conductor.client.http.models.start_workflow_request import StartWorkflowRequest
-from conductor.client.http.rest import ApiException
 from conductor.client.orkes.orkes_workflow_client import OrkesWorkflowClient
 from conductor.client.orkes.orkes_metadata_client import OrkesMetadataClient
 
 from tests.integration.retry_helpers import (
     TERMINAL_WORKFLOW_STATES,
-    is_transient as _is_transient,
     retry_on_transient as _retry_on_transient,
+    wait_for_workflow_terminal,
 )
 
 logging.basicConfig(level=logging.INFO)
@@ -271,22 +270,17 @@ def _wait_for_workflow(self, wf_id, timeout_seconds=POLL_TIMEOUT_SECONDS,
         """Poll until workflow reaches a terminal state. If it doesn't within
         the budget, dump server-side diagnostics so the ensuing assertion shows
         *why* (e.g. a task stuck in SCHEDULED with no poller) rather than only a
-        bare status mismatch.
+        bare status mismatch. Delegates the polling loop to the shared
+        ``wait_for_workflow_terminal`` (transient blips swallowed, real errors
+        raised), then does a definitive final fetch + diagnostics on give-up.
         """
-        deadline = time.monotonic() + timeout_seconds
-        while time.monotonic() < deadline:
-            try:
-                wf = self.workflow_client.get_workflow(wf_id, include_tasks=True)
-            except ApiException as e:
-                # A transient blip on a single poll shouldn't abort the wait;
-                # keep polling until the budget is exhausted.
-                if _is_transient(e):
-                    time.sleep(poll_interval)
-                    continue
-                raise
-            if wf.status in self.TERMINAL_STATES:
-                return wf
-            time.sleep(poll_interval)
+        wf = wait_for_workflow_terminal(
+            self.workflow_client, wf_id,
+            timeout_seconds=timeout_seconds, poll_interval=poll_interval,
+            include_tasks=True, terminal_states=self.TERMINAL_STATES,
+            swallow='transient', log=lambda _msg: None)
+        if wf is not None and wf.status in self.TERMINAL_STATES:
+            return wf
         wf = _retry_on_transient(self.workflow_client.get_workflow,
                                  wf_id, include_tasks=True)
         if wf.status not in self.TERMINAL_STATES:

tests/integration/test_comprehensive_e2e.py

@@ -56,6 +56,7 @@
     TaskExecutionStarted, TaskExecutionCompleted, TaskExecutionFailure,
     TaskUpdateFailure
 )
+from tests.integration.retry_helpers import wait_for_workflow_terminal
 
 # Event collector
 class EventCollector(TaskRunnerEventsListener):
@@ -281,18 +282,24 @@ def _run_workflow_to_terminal(self, workflow_client, timeout_s=90):
         wf_id = workflow_client.start_workflow(start_workflow_request=req)
         print(f"✓ Started workflow: {wf_id}")
 
-        deadline = time.time() + timeout_s
-        wf = None
-        while time.time() < deadline:
-            wf = workflow_client.get_workflow(wf_id, include_tasks=True)
+        # "Terminal" here is stricter than the usual terminal-status check: we
+        # also require the full expected task set to be present, so the caller
+        # never asserts against a half-scheduled workflow (the flaky "4 != 5
+        # tasks" case). Delegates polling to the shared wait_for_workflow_terminal.
+        def _fully_materialized(wf):
+            return getattr(wf, 'status', None) in ('COMPLETED', 'FAILED') \
+                and len(wf.tasks or []) == self.EXPECTED_TASK_COUNT
+
+        def _show(wf):
             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)
 
+        wf = wait_for_workflow_terminal(
+            workflow_client, wf_id,
+            timeout_seconds=timeout_s, poll_interval=1,
+            include_tasks=True, is_terminal=_fully_materialized,
+            swallow='none', log=lambda _msg: None, on_poll=_show)
         return wf_id, wf
 
     def test_03_execute_workflow(self):

tests/integration/test_lease_extension.py

@@ -36,14 +36,13 @@
 from conductor.client.http.models.task_def import TaskDef
 from conductor.client.http.models.workflow_task import WorkflowTask
 from conductor.client.http.models.start_workflow_request import StartWorkflowRequest
-from conductor.client.http.rest import ApiException
 from conductor.client.orkes.orkes_workflow_client import OrkesWorkflowClient
 from conductor.client.orkes.orkes_metadata_client import OrkesMetadataClient
 
 from tests.integration.retry_helpers import (
     TERMINAL_WORKFLOW_STATES,
-    is_transient as _is_transient,
     retry_on_transient as _retry_on_transient,
+    wait_for_workflow_terminal,
 )
 
 logging.basicConfig(level=logging.INFO)
@@ -182,22 +181,17 @@ def _wait_for_workflow(self, wf_id, timeout_seconds=POLL_TIMEOUT_SECONDS,
         """Poll until workflow reaches a terminal state. If it doesn't within
         the budget, dump server-side diagnostics so the ensuing assertion shows
         *why* (e.g. a task stuck in SCHEDULED with no poller) rather than only a
-        bare status mismatch.
+        bare status mismatch. Delegates the polling loop to the shared
+        ``wait_for_workflow_terminal`` (transient blips swallowed, real errors
+        raised), then does a definitive final fetch + diagnostics on give-up.
         """
-        deadline = time.monotonic() + timeout_seconds
-        while time.monotonic() < deadline:
-            try:
-                wf = self.workflow_client.get_workflow(wf_id, include_tasks=True)
-            except ApiException as e:
-                # A transient blip on a single poll shouldn't abort the wait;
-                # keep polling until the budget is exhausted.
-                if _is_transient(e):
-                    time.sleep(poll_interval)
-                    continue
-                raise
-            if wf.status in self.TERMINAL_STATES:
-                return wf
-            time.sleep(poll_interval)
+        wf = wait_for_workflow_terminal(
+            self.workflow_client, wf_id,
+            timeout_seconds=timeout_seconds, poll_interval=poll_interval,
+            include_tasks=True, terminal_states=self.TERMINAL_STATES,
+            swallow='transient', log=lambda _msg: None)
+        if wf is not None and wf.status in self.TERMINAL_STATES:
+            return wf
         wf = _retry_on_transient(self.workflow_client.get_workflow,
                                  wf_id, include_tasks=True)
         if wf.status not in self.TERMINAL_STATES:

tests/integration/workflow/test_workflow_execution.py

@@ -12,7 +12,7 @@
 from conductor.client.workflow.executor.workflow_executor import WorkflowExecutor
 from conductor.client.workflow.task.simple_task import SimpleTask
 from tests.integration.resources.worker.python.python_worker import *
-from tests.integration.retry_helpers import retry_scenario, TERMINAL_WORKFLOW_STATES
+from tests.integration.retry_helpers import retry_scenario, wait_for_workflow_terminal
 
 WORKFLOW_NAME = "sdk_python_integration_test_workflow"
 WORKFLOW_DESCRIPTION = "Python SDK Integration Test"
@@ -207,34 +207,19 @@ def scenario_workflow_execution(
 
 def _wait_for_workflow_terminal(workflow_id, workflow_executor, deadline,
                                 poll_interval=2):
-    """Poll until the workflow reaches a terminal state or the shared deadline
-    passes, logging progress (wf id + elapsed) so a slow-but-eventually-complete
-    run is visible instead of surfacing as a bare timeout. A transient poll error
-    is logged and retried. Returns the last observed status; the caller still
-    runs validate_workflow_status for the actual assertion.
+    """Poll until the workflow reaches a terminal state or the shared ``deadline``
+    (a ``time.time()`` wall-clock value shared across a batch) passes, logging
+    progress so a slow-but-eventually-complete run is visible instead of a bare
+    timeout. A transient poll error is logged and retried. Returns the last
+    observed status; the caller still runs validate_workflow_status for the
+    actual assertion. Thin wrapper over the shared ``wait_for_workflow_terminal``.
     """
-    start = time.time()
-    status = None
-    while True:
-        try:
-            workflow = workflow_executor.get_workflow(
-                workflow_id=workflow_id, include_tasks=False)
-            status = workflow.status
-        except Exception as e:  # transient blip against the shared server
-            logger.warning('error polling workflow %s (%.0fs elapsed): %s',
-                           workflow_id, time.time() - start, e)
-        if status in TERMINAL_WORKFLOW_STATES:
-            logger.info('workflow %s reached %s after %.0fs',
-                        workflow_id, status, time.time() - start)
-            return status
-        if time.time() >= deadline:
-            logger.warning(
-                'workflow %s still %s after %.0fs; giving up wait',
-                workflow_id, status, time.time() - start)
-            return status
-        logger.info('workflow %s still %s after %.0fs; waiting',
-                    workflow_id, status, time.time() - start)
-        sleep(poll_interval)
+    workflow = wait_for_workflow_terminal(
+        workflow_executor, workflow_id,
+        timeout_seconds=max(0.0, deadline - time.time()),
+        poll_interval=poll_interval, include_tasks=False,
+        swallow='all', log=logger.info)
+    return getattr(workflow, 'status', None)
 
 
 def generate_workflow(workflow_executor: WorkflowExecutor, workflow_name: str = WORKFLOW_NAME,
@@ -467,22 +452,13 @@ def scenario_execute_workflow_error_handling(workflow_executor: WorkflowExecutor
 
 
 def _wait_for_workflow_completion(workflow_executor: WorkflowExecutor, workflow_id: str, max_wait_seconds: int = 60):
-    """Helper function to wait for workflow completion"""
-    import time
-    start_time = time.time()
-
-    while time.time() - start_time < max_wait_seconds:
-        workflow = workflow_executor.get_workflow(workflow_id, True)
-
-        if workflow.status in TERMINAL_WORKFLOW_STATES:
-            logger.debug(f'Workflow {workflow_id} finished with status: {workflow.status}')
-            return workflow
-
-        logger.debug(f'Waiting for workflow {workflow_id}... Status: {workflow.status}')
-        time.sleep(2)
-
-    # Return final state even if not completed
-    return workflow_executor.get_workflow(workflow_id, True)
+    """Helper function to wait for workflow completion. Thin wrapper over the
+    shared ``wait_for_workflow_terminal``; returns the last observed Workflow.
+    """
+    return wait_for_workflow_terminal(
+        workflow_executor, workflow_id,
+        timeout_seconds=max_wait_seconds, poll_interval=2,
+        include_tasks=True, swallow='none', log=logger.debug)
 
 # ===== SIGNAL TESTS =====