consolidate retry helper stuff

Author: chrishagglund-ship-it

tests/integration/retry_helpers.py

@@ -1,6 +1,12 @@
-"""Shared transient-retry helper for the aggregate ``test_all`` integration bucket.
+"""Shared transient-retry / flakiness helpers for the integration suites.
 
-The bucket runs against a shared dev server whose transport occasionally blips
+Originally scoped to the aggregate ``test_all`` bucket, this module is now the
+single home for the retry/poll primitives shared across the integration tests:
+``is_transient`` (the one definition of "transient transport blip"),
+``TERMINAL_WORKFLOW_STATES``, the per-request ``retry_on_transient``, and the
+scenario-level ``retry_scenario``.
+
+The suites run against a shared dev server whose transport occasionally blips
 (read timeout, "server disconnected without a response", HTTP/2 GOAWAY, stale
 keep-alive) and surfaces as ``ApiException(status=0)``. Those are not real test
 failures.
@@ -31,6 +37,19 @@
 DEFAULT_BASE_DELAY_SECONDS = 1.0
 DEFAULT_MAX_DELAY_SECONDS = 30.0
 
+# Canonical set of terminal workflow statuses, shared by the poll-to-terminal
+# helpers across the integration suites so the set can't drift file to file.
+TERMINAL_WORKFLOW_STATES = ('COMPLETED', 'FAILED', 'TIMED_OUT', 'TERMINATED')
+
+
+def is_transient(exc):
+    """True when ``exc`` is a transient transport blip against the shared dev
+    server (read timeout, connection reset, HTTP/2 GOAWAY, stale keep-alive)
+    rather than a real failure. status 0/None means no HTTP response arrived.
+    """
+    return isinstance(exc, ApiException) and (
+        getattr(exc, 'transient', False) or exc.status in (0, None))
+
 
 def first_transient_api_exception(exc):
     """Walk the exception chain (``__cause__`` / ``__context__``) and return the
@@ -46,13 +65,34 @@ def first_transient_api_exception(exc):
     cur = exc
     while cur is not None and id(cur) not in seen:
         seen.add(id(cur))
-        if isinstance(cur, ApiException) and (
-                getattr(cur, 'transient', False) or cur.status in (0, None)):
+        if is_transient(cur):
             return cur
         cur = cur.__cause__ or cur.__context__
     return None
 
 
+def retry_on_transient(func, *args, retries=4, base_delay=1, **kwargs):
+    """Retry ``func(*args, **kwargs)`` on a transient (status 0) transport blip
+    with capped-free exponential backoff. Non-transient errors (real 4xx/5xx,
+    assertion failures) raise immediately.
+
+    This is *per-request* retry, suited to idempotent calls (get/register).
+    For non-idempotent scenarios (start_workflow, signal) use ``retry_scenario``,
+    which re-runs the whole scenario from scratch instead.
+    """
+    for attempt in range(retries):
+        try:
+            return func(*args, **kwargs)
+        except ApiException as e:
+            if is_transient(e) and attempt < retries - 1:
+                logger.warning(
+                    'transient (%s) API error (attempt %d/%d): %s; retrying',
+                    e.status, attempt + 1, retries, e)
+                time.sleep(base_delay * (2 ** attempt))
+                continue
+            raise
+
+
 def retry_scenario(label, func, *args, deadline=None,
                    base_delay=DEFAULT_BASE_DELAY_SECONDS,
                    max_delay=DEFAULT_MAX_DELAY_SECONDS, **kwargs):

tests/integration/test_async_lease_extension.py

@@ -42,35 +42,15 @@
 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,
+)
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-
-def _is_transient(e):
-    """A transient (status 0) ApiException is a raw transport blip against the
-    shared dev server (read timeout, connection reset, HTTP/2 GOAWAY, stale
-    keep-alive) — not a real failure. status 0 means no HTTP response was
-    received at all."""
-    return isinstance(e, ApiException) and (
-        getattr(e, 'transient', False) or e.status in (0, None))
-
-
-def _retry_on_transient(func, *args, retries=4, base_delay=1, **kwargs):
-    """Retry a workflow/metadata client call on a transient (0) transport blip.
-    These `(0)` errors flake the lease suites on a loaded shared server; retrying
-    absorbs the noise. Non-transient errors (real 4xx/5xx) raise immediately."""
-    for attempt in range(retries):
-        try:
-            return func(*args, **kwargs)
-        except ApiException as e:
-            if _is_transient(e) and attempt < retries - 1:
-                logger.warning(
-                    'transient (%s) API error (attempt %d/%d): %s; retrying',
-                    e.status, attempt + 1, retries, e)
-                time.sleep(base_delay * (2 ** attempt))
-                continue
-            raise
-
 # Short response timeout — task must heartbeat to stay alive
 RESPONSE_TIMEOUT_SECONDS = 10
 
@@ -284,7 +264,7 @@ def _start_workflow(self, wf_name, job_id):
         logger.info("Started workflow %s: %s", wf_name, wf_id)
         return wf_id
 
-    TERMINAL_STATES = ('COMPLETED', 'FAILED', 'TIMED_OUT', 'TERMINATED')
+    TERMINAL_STATES = TERMINAL_WORKFLOW_STATES
 
     def _wait_for_workflow(self, wf_id, timeout_seconds=POLL_TIMEOUT_SECONDS,
                            poll_interval=POLL_INTERVAL_SECONDS):

tests/integration/test_lease_extension.py

@@ -40,35 +40,16 @@
 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,
+)
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-def _is_transient(e):
-    """A transient (status 0) ApiException is a raw transport blip against the
-    shared dev server (read timeout, connection reset, HTTP/2 GOAWAY, stale
-    keep-alive) — not a real failure. status 0 means no HTTP response was
-    received at all."""
-    return isinstance(e, ApiException) and (
-        getattr(e, 'transient', False) or e.status in (0, None))
-
-
-def _retry_on_transient(func, *args, retries=4, base_delay=1, **kwargs):
-    """Retry a workflow/metadata client call on a transient (0) transport blip.
-    These `(0)` errors flake the lease suites on a loaded shared server; retrying
-    absorbs the noise. Non-transient errors (real 4xx/5xx) raise immediately."""
-    for attempt in range(retries):
-        try:
-            return func(*args, **kwargs)
-        except ApiException as e:
-            if _is_transient(e) and attempt < retries - 1:
-                logger.warning(
-                    'transient (%s) API error (attempt %d/%d): %s; retrying',
-                    e.status, attempt + 1, retries, e)
-                time.sleep(base_delay * (2 ** attempt))
-                continue
-            raise
-
 # Short response timeout — task must heartbeat to stay alive
 RESPONSE_TIMEOUT_SECONDS = 10
 
@@ -194,7 +175,7 @@ def _start_workflow(self, wf_name, job_id):
         logger.info("Started workflow %s: %s", wf_name, wf_id)
         return wf_id
 
-    TERMINAL_STATES = ('COMPLETED', 'FAILED', 'TIMED_OUT', 'TERMINATED')
+    TERMINAL_STATES = TERMINAL_WORKFLOW_STATES
 
     def _wait_for_workflow(self, wf_id, timeout_seconds=POLL_TIMEOUT_SECONDS,
                            poll_interval=POLL_INTERVAL_SECONDS):

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
+from tests.integration.retry_helpers import retry_scenario, TERMINAL_WORKFLOW_STATES
 
 WORKFLOW_NAME = "sdk_python_integration_test_workflow"
 WORKFLOW_DESCRIPTION = "Python SDK Integration Test"
@@ -213,7 +213,6 @@ def _wait_for_workflow_terminal(workflow_id, workflow_executor, deadline,
     is logged and retried. Returns the last observed status; the caller still
     runs validate_workflow_status for the actual assertion.
     """
-    terminal = ('COMPLETED', 'FAILED', 'TIMED_OUT', 'TERMINATED')
     start = time.time()
     status = None
     while True:
@@ -224,7 +223,7 @@ def _wait_for_workflow_terminal(workflow_id, workflow_executor, deadline,
         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:
+        if status in TERMINAL_WORKFLOW_STATES:
             logger.info('workflow %s reached %s after %.0fs',
                         workflow_id, status, time.time() - start)
             return status
@@ -475,7 +474,7 @@ def _wait_for_workflow_completion(workflow_executor: WorkflowExecutor, workflow_
     while time.time() - start_time < max_wait_seconds:
         workflow = workflow_executor.get_workflow(workflow_id, True)
 
-        if workflow.status in ['COMPLETED', 'FAILED', 'TERMINATED', 'TIMED_OUT']:
+        if workflow.status in TERMINAL_WORKFLOW_STATES:
             logger.debug(f'Workflow {workflow_id} finished with status: {workflow.status}')
             return workflow
 
@@ -888,22 +887,4 @@ def scenario_signal_to_dict_fix(workflow_executor: WorkflowExecutor):
 
     _wait_for_workflow_completion(workflow_executor, workflow_id)
 
-    logger.info('to_dict() method test completed')
-
-
-def _wait_for_workflow_completion(workflow_executor: WorkflowExecutor, workflow_id: str, timeout: int = 10):
-    """Wait for workflow to complete with timeout"""
-    max_iterations = timeout * 10  # Check every 0.1 seconds
-
-    for i in range(max_iterations):
-        try:
-            workflow = workflow_executor.get_workflow(workflow_id, include_tasks=False)
-            if workflow.status in ["COMPLETED", "FAILED", "TERMINATED"]:
-                logger.debug(f'Workflow {workflow_id} completed with status: {workflow.status}')
-                return
-        except Exception as e:
-            logger.warning(f'Error checking workflow status: {e}')
-
-        time.sleep(0.1)
-
-    logger.warning(f'Workflow {workflow_id} did not complete within {timeout} seconds')
+    logger.info('to_dict() method test completed')