[airflow] branch main updated: Respect `soft_fail` parameter in `LambdaFunctionStateSensor` (#34551)

[vincbeck]

This is an automated email from the ASF dual-hosted git repository.

vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new b4feb5c20d Respect `soft_fail` parameter in 
`LambdaFunctionStateSensor` (#34551)
b4feb5c20d is described below

commit b4feb5c20d463162897b6a13a3c27f176c56fd99
Author: Utkarsh Sharma <utkarshar...@gmail.com>
AuthorDate: Fri Sep 22 20:47:32 2023 +0530

    Respect `soft_fail` parameter in `LambdaFunctionStateSensor` (#34551)
---
 .../providers/amazon/aws/sensors/lambda_function.py    | 10 ++++++----
 .../amazon/aws/sensors/test_lambda_function.py         | 18 +++++++++++++++++-
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/airflow/providers/amazon/aws/sensors/lambda_function.py 
b/airflow/providers/amazon/aws/sensors/lambda_function.py
index 772ed0689a..579c49b76d 100644
--- a/airflow/providers/amazon/aws/sensors/lambda_function.py
+++ b/airflow/providers/amazon/aws/sensors/lambda_function.py
@@ -26,7 +26,7 @@ from airflow.providers.amazon.aws.utils import 
trim_none_values
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.sensors.base import BaseSensorOperator
 
 
@@ -74,9 +74,11 @@ class LambdaFunctionStateSensor(BaseSensorOperator):
         state = 
self.hook.conn.get_function(**trim_none_values(get_function_args))["Configuration"]["State"]
 
         if state in self.FAILURE_STATES:
-            raise AirflowException(
-                "Lambda function state sensor failed because the Lambda is in 
a failed state"
-            )
+            message = "Lambda function state sensor failed because the Lambda 
is in a failed state"
+            # TODO: remove this if block when min_airflow_version is set to 
higher than 2.7.1
+            if self.soft_fail:
+                raise AirflowSkipException(message)
+            raise AirflowException(message)
 
         return state in self.target_states
 
diff --git a/tests/providers/amazon/aws/sensors/test_lambda_function.py 
b/tests/providers/amazon/aws/sensors/test_lambda_function.py
index f97eb68b34..652d8996fd 100644
--- a/tests/providers/amazon/aws/sensors/test_lambda_function.py
+++ b/tests/providers/amazon/aws/sensors/test_lambda_function.py
@@ -20,7 +20,7 @@ from unittest import mock
 
 import pytest
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
 from airflow.providers.amazon.aws.sensors.lambda_function import 
LambdaFunctionStateSensor
 
@@ -69,3 +69,19 @@ class TestLambdaFunctionStateSensor:
             mock_conn.get_function.assert_called_once_with(
                 FunctionName=FUNCTION_NAME,
             )
+
+    @pytest.mark.parametrize(
+        "soft_fail, expected_exception", ((False, AirflowException), (True, 
AirflowSkipException))
+    )
+    def test_fail_poke(self, soft_fail, expected_exception):
+        sensor = LambdaFunctionStateSensor(
+            task_id="test_sensor",
+            function_name=FUNCTION_NAME,
+        )
+        sensor.soft_fail = soft_fail
+        message = "Lambda function state sensor failed because the Lambda is 
in a failed state"
+        with pytest.raises(expected_exception, match=message), mock.patch(
+            
"airflow.providers.amazon.aws.hooks.lambda_function.LambdaHook.conn"
+        ) as conn:
+            conn.get_function.return_value = {"Configuration": {"State": 
"Failed"}}
+            sensor.poke(context={})