additional retry for stuff that is seeing transient errors/flakiness

Author: chrishagglund-ship-it

tests/integration/metadata/test_schema_service.py

@@ -5,6 +5,7 @@
 from conductor.client.http.api.schema_resource_api import SchemaResourceApi
 from conductor.client.http.models.schema_def import SchemaDef, SchemaType
 from conductor.client.orkes.orkes_schema_client import OrkesSchemaClient
+from tests.integration.retry_helpers import retry_on_transient, retry_on_status
 
 SCHEMA_NAME = 'ut_schema'
 SCHEMA_VERSION = 1
@@ -41,33 +42,40 @@ def test_init(self):
         self.assertIsInstance(self.schema_client.schemaApi, SchemaResourceApi, message)
 
     def test_registerSchema(self):
-        self.schema_client.register_schema(self.schemaDef)
-        response = self.schema_client.schemaApi.get_schema_by_name_and_version(name=SCHEMA_NAME, version=SCHEMA_VERSION)
+        retry_on_transient(self.schema_client.register_schema, self.schemaDef)
+        # A GET right after register can briefly 404 until the write propagates
+        # on the shared dev server; retry the read rather than fail the test.
+        response = retry_on_status(
+            self.schema_client.schemaApi.get_schema_by_name_and_version,
+            name=SCHEMA_NAME, version=SCHEMA_VERSION)
         self.assertEqual(response.name, SCHEMA_NAME)
         self.assertEqual(response.version, SCHEMA_VERSION)
         self.assertEqual(response.type, SchemaType.JSON)
 
     def test_getSchema(self):
-        self.schema_client.register_schema(self.schemaDef)
-        schema = self.schema_client.get_schema(SCHEMA_NAME, SCHEMA_VERSION)
+        retry_on_transient(self.schema_client.register_schema, self.schemaDef)
+        # A GET right after register can briefly 404 until the write propagates
+        # on the shared dev server; retry the read rather than fail the test.
+        schema = retry_on_status(self.schema_client.get_schema,
+                                 SCHEMA_NAME, SCHEMA_VERSION)
         self.assertEqual(schema.name, SCHEMA_NAME)
         self.assertEqual(schema.version, SCHEMA_VERSION)
 
     def test_getAllSchemas(self):
         schemaDef2 = SchemaDef(name='ut_schema_2', version=1, type=SchemaType.JSON, data=schema, external_ref='http://example.com/2')
-        self.schema_client.register_schema(self.schemaDef)
-        self.schema_client.register_schema(schemaDef2)
+        retry_on_transient(self.schema_client.register_schema, self.schemaDef)
+        retry_on_transient(self.schema_client.register_schema, schemaDef2)
         schemas = self.schema_client.get_all_schemas()
         self.assertGreaterEqual(len(schemas), 2)
 
     def test_deleteSchema(self):
-        self.schema_client.register_schema(self.schemaDef)
+        retry_on_transient(self.schema_client.register_schema, self.schemaDef)
         self.schema_client.delete_schema(SCHEMA_NAME, SCHEMA_VERSION)
         with self.assertRaises(Exception):
             self.schema_client.get_schema(SCHEMA_NAME, SCHEMA_VERSION)
 
     def test_deleteSchemaByName(self):
-        self.schema_client.register_schema(self.schemaDef)
+        retry_on_transient(self.schema_client.register_schema, self.schemaDef)
         self.schema_client.delete_schema_by_name(SCHEMA_NAME)
         with self.assertRaises(Exception):
             self.schema_client.get_schema(SCHEMA_NAME, SCHEMA_VERSION)

tests/integration/retry_helpers.py

@@ -93,6 +93,34 @@ def retry_on_transient(func, *args, retries=4, base_delay=1, **kwargs):
             raise
 
 
+def retry_on_status(func, *args, statuses=(404,), retries=5, base_delay=1.0,
+                    max_delay=None, **kwargs):
+    """Retry ``func(*args, **kwargs)`` on an ``ApiException`` whose status is in
+    ``statuses`` (in addition to transient status-0 blips), with capped
+    exponential backoff. Any other error raises immediately.
+
+    Intended for read-after-write races against the shared dev server: a GET
+    issued right after a register/update can briefly 404 until the write
+    propagates. This is *per-request* retry, so use it only for idempotent reads
+    (or writes that are safe to repeat).
+    """
+    for attempt in range(retries):
+        try:
+            return func(*args, **kwargs)
+        except ApiException as e:
+            retryable = is_transient(e) or e.status in statuses
+            if retryable and attempt < retries - 1:
+                delay = base_delay * (2 ** attempt)
+                if max_delay is not None:
+                    delay = min(delay, max_delay)
+                logger.warning(
+                    'retryable (%s) API error (attempt %d/%d): %s; retrying in '
+                    '%.1fs', e.status, attempt + 1, retries, e, delay)
+                time.sleep(delay)
+                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_v2_fallback_intg.py

@@ -25,6 +25,7 @@
 from conductor.client.orkes.orkes_metadata_client import OrkesMetadataClient
 from conductor.client.orkes.orkes_workflow_client import OrkesWorkflowClient
 from conductor.client.worker.worker_task import worker_task
+from tests.integration.retry_helpers import retry_on_transient
 
 logger = logging.getLogger(__name__)
 
@@ -84,10 +85,14 @@ def test_0_register_workflow(self):
 
         workflow = WorkflowDef(name=WORKFLOW_NAME, version=WORKFLOW_VERSION)
         workflow._tasks = tasks
+        # Retry registration on a transient (status 0) transport blip against the
+        # shared dev server so a dropped connection doesn't fail the suite.
         try:
-            self.metadata_client.update_workflow_def(workflow, overwrite=True)
+            retry_on_transient(self.metadata_client.update_workflow_def,
+                               workflow, overwrite=True)
         except Exception:
-            self.metadata_client.register_workflow_def(workflow, overwrite=True)
+            retry_on_transient(self.metadata_client.register_workflow_def,
+                               workflow, overwrite=True)
         print(f"\n  Registered workflow '{WORKFLOW_NAME}' with {len(tasks)} tasks")
 
     def test_1_workflows_complete_with_v2_or_fallback(self):
@@ -128,7 +133,11 @@ def _run_workers():
                 req.name = WORKFLOW_NAME
                 req.version = WORKFLOW_VERSION
                 req.input = {"run_index": i}
-                wf_id = self.workflow_client.start_workflow(start_workflow_request=req)
+                # A status-0 blip here means no response arrived, so no id was
+                # returned; retrying gets a fresh attempt (any orphaned run just
+                # completes untracked and doesn't affect the tracked-id count).
+                wf_id = retry_on_transient(self.workflow_client.start_workflow,
+                                           start_workflow_request=req)
                 workflow_ids.append(wf_id)
 
             print(f"\n  Submitted {len(workflow_ids)} workflows")