From 9d48300e90373a053e1fd6dd53689375ebbe22ef Mon Sep 17 00:00:00 2001 From: beobest2 Date: Thu, 16 May 2024 01:06:16 -0400 Subject: [PATCH] fix test code --- .../amazon/aws/operators/test_emr_containers.py | 7 +++++-- .../amazon/aws/sensors/test_emr_containers.py | 16 +++++++++------- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/tests/providers/amazon/aws/operators/test_emr_containers.py b/tests/providers/amazon/aws/operators/test_emr_containers.py index fa5c4764b87ea8..feeec1278e1552 100644 --- a/tests/providers/amazon/aws/operators/test_emr_containers.py +++ b/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: diff --git a/tests/providers/amazon/aws/sensors/test_emr_containers.py b/tests/providers/amazon/aws/sensors/test_emr_containers.py index 048533cb60e6c6..65ae0729341ff0 100644 --- a/tests/providers/amazon/aws/sensors/test_emr_containers.py +++ b/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