Skip to content

Commit ae55e22

Browse files
authored
Make azure test less flaky/racy (apache#46281)
1 parent d4d3e54 commit ae55e22

1 file changed

Lines changed: 20 additions & 14 deletions

File tree

providers/tests/microsoft/azure/operators/test_container_instances.py

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -453,24 +453,30 @@ def test_execute_fails_with_incorrect_restart_policy(self, aci_mock):
453453
)
454454

455455
@mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook")
456-
@mock.patch("airflow.providers.microsoft.azure.operators.container_instances.time.sleep")
457-
def test_execute_correct_sleep_cycle(self, sleep_mock, aci_mock):
456+
def test_execute_correct_sleep_cycle(self, aci_mock):
458457
expected_cg1 = make_mock_container(state="Running", exit_code=0, detail_status="test")
459458
expected_cg2 = make_mock_container(state="Terminated", exit_code=0, detail_status="test")
460459

461-
aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg1, expected_cg2]
462-
aci_mock.return_value.exists.return_value = False
460+
import time
463461

464-
aci = AzureContainerInstancesOperator(
465-
ci_conn_id=None,
466-
registry_conn_id=None,
467-
resource_group="resource-group",
468-
name="container-name",
469-
image="container-image",
470-
region="region",
471-
task_id="task",
472-
)
473-
aci.execute(None)
462+
original_time_sleep = time.sleep
463+
with mock.patch(
464+
"airflow.providers.microsoft.azure.operators.container_instances.time.sleep"
465+
) as sleep_mock:
466+
sleep_mock.side_effect = lambda _: original_time_sleep(0.1)
467+
aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg1, expected_cg2]
468+
aci_mock.return_value.exists.return_value = False
469+
470+
aci = AzureContainerInstancesOperator(
471+
ci_conn_id=None,
472+
registry_conn_id=None,
473+
resource_group="resource-group",
474+
name="container-name",
475+
image="container-image",
476+
region="region",
477+
task_id="task",
478+
)
479+
aci.execute(None)
474480

475481
# sleep is called at the end of cycles. Thus, the Terminated call does not trigger sleep
476482
assert sleep_mock.call_count == 2

0 commit comments

Comments
 (0)