@@ -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