fix test code

tests/providers/amazon/aws/operators/test_emr_containers.py

@@ -152,10 +152,13 @@ def test_operator_defer_with_timeout(self, mock_submit_job, mock_check_query_sta
  self.emr_container.deferrable = True
  self.emr_container.max_polling_attempts = 1000
 
- with pytest.raises(TaskDeferred):
+ with pytest.raises(TaskDeferred) as e:
  self.emr_container.execute(context=None)
 
- assert mock_check_query_status.call_count == 1000
+ trigger = e.value.trigger
+ assert isinstance(trigger, EmrContainerTrigger), f"{trigger} is not a EmrContainerTrigger"
+ assert trigger.waiter_delay == self.emr_container.poll_interval
+ assert trigger.attempts == self.emr_container.max_polling_attempts
 
 
 class TestEmrEksCreateClusterOperator:

tests/providers/amazon/aws/sensors/test_emr_containers.py

@@ -85,14 +85,16 @@ def test_sensor_defer(self, mock_poke):
  e.value.trigger, EmrContainerTrigger
  ), f"{e.value.trigger} is not a EmrContainerTrigger"
 
- @mock.patch.object(
- EmrContainerHook, "check_query_status", return_value=EmrContainerHook.INTERMEDIATE_STATES[0]
- )
- def test_sensor_defer_with_timeout(self, mock_check_query_status):
+ @mock.patch("airflow.providers.amazon.aws.sensors.emr.EmrContainerSensor.poke")
+ def test_sensor_defer_with_timeout(self, mock_poke):
  self.sensor.deferrable = True
- self.sensor.max_polling_attempts = 1000
+ mock_poke.return_value = False
+ self.sensor.max_retries = 1000
 
- with pytest.raises(TaskDeferred):
+ with pytest.raises(TaskDeferred) as e:
  self.sensor.execute(context=None)
 
- assert mock_check_query_status.call_count == 1000
+ trigger = e.value.trigger
+ assert isinstance(trigger, EmrContainerTrigger), f"{trigger} is not a EmrContainerTrigger"
+ assert trigger.waiter_delay == self.sensor.poll_interval
+ assert trigger.attempts == self.sensor.max_retries