From 9b2bb19914ea8f2b7301c5abec149cd723fe29ce Mon Sep 17 00:00:00 2001 From: bingqin2 Date: Wed, 9 Sep 2026 11:18:48 -0500 Subject: [PATCH] Add failure-path and serialization tests for StepFunctionsExecutionCompleteTrigger The region_name fix originally proposed in this PR landed via #72625. This adds the remaining coverage: the error TriggerEvent path, verify/botocore_config serialization, pruning of unset hook parameters, and the region_name docstring. --- .../amazon/aws/triggers/step_function.py | 2 + .../amazon/aws/triggers/test_step_function.py | 41 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/step_function.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/step_function.py index a52f190969643..04e22f924bce6 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/step_function.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/step_function.py @@ -33,6 +33,8 @@ class StepFunctionsExecutionCompleteTrigger(AwsBaseWaiterTrigger): :param waiter_delay: The amount of time in seconds to wait between attempts. :param waiter_max_attempts: The maximum number of attempts to be made. :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: AWS region name to use. + Override the region_name in connection (if provided). """ def __init__( diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_step_function.py b/providers/amazon/tests/unit/amazon/aws/triggers/test_step_function.py index 3341952450eb8..4ee0cc4555b41 100644 --- a/providers/amazon/tests/unit/amazon/aws/triggers/test_step_function.py +++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_step_function.py @@ -21,6 +21,7 @@ import pytest +from airflow.exceptions import AirflowException from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook from airflow.providers.amazon.aws.triggers.step_function import StepFunctionsExecutionCompleteTrigger from airflow.triggers.base import TriggerEvent @@ -87,3 +88,43 @@ async def test_run_success(self, mock_async_conn, mock_get_waiter): self.EXPECTED_WAITER_NAME, deferrable=True, client=mock.ANY, config_overrides=None ) assert mock_get_waiter().wait.call_args.kwargs["executionArn"] == self.EXECUTION_ARN + + def test_serialization_with_verify_and_botocore_config(self): + trigger = StepFunctionsExecutionCompleteTrigger( + execution_arn=self.EXECUTION_ARN, + aws_conn_id="aws_step_function_conn", + region_name="eu-central-1", + verify=False, + botocore_config={"connect_timeout": 30}, + ) + + classpath, kwargs = trigger.serialize() + + assert classpath == BASE_TRIGGER_CLASSPATH + "StepFunctionsExecutionCompleteTrigger" + assert kwargs["verify"] is False + assert kwargs["botocore_config"] == {"connect_timeout": 30} + + def test_serialization_omits_unset_hook_params(self): + trigger = StepFunctionsExecutionCompleteTrigger(execution_arn=self.EXECUTION_ARN) + + _, kwargs = trigger.serialize() + + assert "region_name" not in kwargs + assert "verify" not in kwargs + assert "botocore_config" not in kwargs + + @pytest.mark.asyncio + @mock.patch("airflow.providers.amazon.aws.triggers.base.async_wait") + @mock.patch.object(StepFunctionHook, "get_waiter") + @mock.patch.object(StepFunctionHook, "get_async_conn") + async def test_run_failure(self, mock_async_conn, mock_get_waiter, mock_async_wait): + mock_async_conn.return_value.__aenter__.return_value = mock.MagicMock() + mock_async_wait.side_effect = AirflowException("Step function failed") + trigger = StepFunctionsExecutionCompleteTrigger(execution_arn=self.EXECUTION_ARN) + + generator = trigger.run() + response = await generator.asend(None) + + assert response == TriggerEvent( + {"status": "error", "message": "Step function failed", "execution_arn": self.EXECUTION_ARN} + )