Make azure test less flaky/racy (apache#46281)

ambika-garg · Feb 13, 2025 · 5d3aae0 · 5d3aae0
1 parent 388d47f
commit 5d3aae0
Showing 1 changed file with 20 additions and 14 deletions.
diff --git a/providers/tests/microsoft/azure/operators/test_container_instances.py b/providers/tests/microsoft/azure/operators/test_container_instances.py
@@ -453,24 +453,30 @@ def test_execute_fails_with_incorrect_restart_policy(self, aci_mock):
         )
 
     @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook")
-    @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.time.sleep")
-    def test_execute_correct_sleep_cycle(self, sleep_mock, aci_mock):
+    def test_execute_correct_sleep_cycle(self, aci_mock):
         expected_cg1 = make_mock_container(state="Running", exit_code=0, detail_status="test")
         expected_cg2 = make_mock_container(state="Terminated", exit_code=0, detail_status="test")
 
-        aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg1, expected_cg2]
-        aci_mock.return_value.exists.return_value = False
+        import time
 
-        aci = AzureContainerInstancesOperator(
-            ci_conn_id=None,
-            registry_conn_id=None,
-            resource_group="resource-group",
-            name="container-name",
-            image="container-image",
-            region="region",
-            task_id="task",
-        )
-        aci.execute(None)
+        original_time_sleep = time.sleep
+        with mock.patch(
+            "airflow.providers.microsoft.azure.operators.container_instances.time.sleep"
+        ) as sleep_mock:
+            sleep_mock.side_effect = lambda _: original_time_sleep(0.1)
+            aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg1, expected_cg2]
+            aci_mock.return_value.exists.return_value = False
+
+            aci = AzureContainerInstancesOperator(
+                ci_conn_id=None,
+                registry_conn_id=None,
+                resource_group="resource-group",
+                name="container-name",
+                image="container-image",
+                region="region",
+                task_id="task",
+            )
+            aci.execute(None)
 
         # sleep is called at the end of cycles. Thus, the Terminated call does not trigger sleep
         assert sleep_mock.call_count == 2