From 69a74200ea4feb1f46385b2bf061d1a5e904d457 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Sep 2023 19:14:05 +0530 Subject: [PATCH] Respect soft_fail parameter in StepFunctionExecutionSensor --- .../amazon/aws/sensors/step_function.py | 8 ++++++-- .../amazon/aws/sensors/test_step_function.py | 18 +++++++++++++++++- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/airflow/providers/amazon/aws/sensors/step_function.py b/airflow/providers/amazon/aws/sensors/step_function.py index 2b0c3b607354e..053a751336268 100644 --- a/airflow/providers/amazon/aws/sensors/step_function.py +++ b/airflow/providers/amazon/aws/sensors/step_function.py @@ -22,7 +22,7 @@ from deprecated import deprecated -from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning +from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowSkipException from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook from airflow.sensors.base import BaseSensorOperator @@ -76,7 +76,11 @@ def poke(self, context: Context): output = json.loads(execution_status["output"]) if "output" in execution_status else None if state in self.FAILURE_STATES: - raise AirflowException(f"Step Function sensor failed. State Machine Output: {output}") + # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1 + message = f"Step Function sensor failed. State Machine Output: {output}" + if self.soft_fail: + raise AirflowSkipException(message) + raise AirflowException(message) if state in self.INTERMEDIATE_STATES: return False diff --git a/tests/providers/amazon/aws/sensors/test_step_function.py b/tests/providers/amazon/aws/sensors/test_step_function.py index d0452dd0c2810..b6a47d49cd642 100644 --- a/tests/providers/amazon/aws/sensors/test_step_function.py +++ b/tests/providers/amazon/aws/sensors/test_step_function.py @@ -17,12 +17,13 @@ # under the License. from __future__ import annotations +import json from unittest import mock from unittest.mock import MagicMock import pytest -from airflow.exceptions import AirflowException +from airflow.exceptions import AirflowException, AirflowSkipException from airflow.providers.amazon.aws.sensors.step_function import StepFunctionExecutionSensor TASK_ID = "step_function_execution_sensor" @@ -88,3 +89,18 @@ def test_succeeded(self, mock_hook): ) assert sensor.poke(self.mock_context) + + @pytest.mark.parametrize( + "soft_fail, expected_exception", ((False, AirflowException), (True, AirflowSkipException)) + ) + @mock.patch("airflow.providers.amazon.aws.hooks.step_function.StepFunctionHook.describe_execution") + def test_fail_poke(self, describe_execution, soft_fail, expected_exception): + sensor = StepFunctionExecutionSensor( + task_id=TASK_ID, execution_arn=EXECUTION_ARN, aws_conn_id=AWS_CONN_ID, region_name=REGION_NAME + ) + sensor.soft_fail = soft_fail + output = '{"test": "test"}' + describe_execution.return_value = {"status": "FAILED", "output": output} + message = f"Step Function sensor failed. State Machine Output: {json.loads(output)}" + with pytest.raises(expected_exception, match=message): + sensor.poke(context={})