From 487d5b9a06d037264c922a4fb8f32a800550009e Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Wed, 26 Aug 2026 14:52:31 +0530 Subject: [PATCH 1/2] Use the operator's AWS settings for deferred Neptune, MWAA, SSM tasks An AwsBaseOperator/AwsBaseSensor subclass resolves region_name, verify and botocore_config in __init__, but did not hand them to the trigger it defers to. The trigger builds its own hook, so the deferred half of the task reached AWS with the default region, SSL verification silently re-enabled, and any custom botocore timeouts or retries discarded. The triggers already accept all three, so only the call sites were missing. --- .../providers/amazon/aws/operators/mwaa.py | 3 + .../amazon/aws/operators/neptune_analytics.py | 24 ++ .../providers/amazon/aws/sensors/mwaa.py | 6 + .../providers/amazon/aws/sensors/ssm.py | 3 + .../unit/amazon/aws/operators/test_mwaa.py | 21 ++ .../aws/operators/test_neptune_analytics.py | 205 ++++++++++++++++++ .../unit/amazon/aws/sensors/test_mwaa.py | 38 +++- .../tests/unit/amazon/aws/sensors/test_ssm.py | 20 ++ 8 files changed, 319 insertions(+), 1 deletion(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py index d9cd6cbdf38d6..ee468f7a4cdfd 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py @@ -170,6 +170,9 @@ def execute(self, context: Context) -> dict: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py index 134d13df982c5..239906e5e9358 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py @@ -173,6 +173,9 @@ def execute(self, context: Context) -> dict: self.defer( trigger=NeptuneGraphAvailableTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_id, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -317,6 +320,9 @@ def execute(self, context: Context) -> dict: self.defer( trigger=NeptuneGraphPrivateEndpointAvailableTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_identifier, vpc_id=self.vpc_id, waiter_delay=self.waiter_delay, @@ -420,6 +426,9 @@ def execute(self, context: Context) -> None: self.defer( trigger=NeptuneGraphPrivateEndpointDeletedTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_identifier, vpc_id=self.vpc_id, endpoint_id=endpoint_id, @@ -517,6 +526,9 @@ def execute(self, context: Context): self.defer( trigger=NeptuneGraphDeletedTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_id, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -729,6 +741,9 @@ def execute(self, context: Context) -> dict: self.defer( trigger=NeptuneGraphAvailableTrigger( aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, graph_id=self.graph_id, waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, @@ -773,6 +788,9 @@ def defer_wait_for_task( waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", kwargs={"graph_id": graph_id}, @@ -914,6 +932,9 @@ def execute(self, context: Context) -> dict: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -1002,6 +1023,9 @@ def execute(self, context: Context) -> dict: waiter_delay=self.waiter_delay, waiter_max_attempts=self.waiter_max_attempts, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py index 37710a7f0bf0c..fb0fe4ea5df7d 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py @@ -161,6 +161,9 @@ def execute(self, context: Context): waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) @@ -312,6 +315,9 @@ def execute(self, context: Context): waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, ), method_name="execute_complete", ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/ssm.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/ssm.py index 943c960873095..d4544c7b7369f 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/ssm.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/ssm.py @@ -142,6 +142,9 @@ def execute(self, context: Context): waiter_delay=int(self.poke_interval), waiter_max_attempts=self.max_retries, aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, fail_on_nonzero_exit=self.fail_on_nonzero_exit, ), method_name="execute_complete", diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py index 566b4bee61384..e2232f77384fd 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py @@ -23,6 +23,7 @@ from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook from airflow.providers.amazon.aws.operators.mwaa import MwaaTriggerDagRunOperator +from airflow.providers.common.compat.sdk import TaskDeferred from unit.amazon.aws.utils.test_template_fields import validate_template_fields @@ -41,6 +42,9 @@ "waiter_max_attempts": 20, "deferrable": False, } +REGION_NAME = "eu-west-2" +VERIFY = False +BOTOCORE_CONFIG = {"read_timeout": 42} HOOK_RETURN_VALUE = { "ResponseMetadata": {}, "RestApiStatusCode": 200, @@ -115,3 +119,20 @@ def test_execute_wait_combinations(self, mock_hook, _, wait_for_completion, defe assert response == HOOK_RETURN_VALUE assert mock_hook.get_waiter.call_count == wait_for_completion assert op.defer.call_count == deferrable + + @mock.patch.object(MwaaTriggerDagRunOperator, "hook") + def test_deferred_trigger_receives_hook_configuration(self, mock_hook): + mock_hook.invoke_rest_api.return_value = HOOK_RETURN_VALUE + op = MwaaTriggerDagRunOperator( + **{**OP_KWARGS, "wait_for_completion": False, "deferrable": True}, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + op.execute({}) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py b/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py index 472e386888c31..f5aea2c2fc95a 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py @@ -52,6 +52,9 @@ ENDPOINT_ID = "vpce-12345" SOURCE_S3_URI = "s3://my-bucket/my-data/" ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole" +REGION_NAME = "eu-west-2" +VERIFY = False +BOTOCORE_CONFIG = {"read_timeout": 42} class TestNeptuneCreateGraphOperator: @@ -219,6 +222,29 @@ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_pe assert isinstance(trigger, NeptuneGraphAvailableTrigger) assert exc_info.value.method_name == "execute_complete" + @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist") + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_persist): + mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"} + + operator = NeptuneCreateGraphOperator( + task_id="test_task", + graph_name=GRAPH_NAME, + vector_search_config={"test": 123}, + provisioned_memory=16, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneCreatePrivateGraphEndpointOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -417,6 +443,33 @@ def test_execute_complete(self, mock_get_endpoint, mock_conn): ) assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID} + @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.VpcEndpointLink.persist") + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_persist): + mock_conn.create_private_graph_endpoint.return_value = { + "status": "CREATING", + "vpcEndpointId": ENDPOINT_ID, + "vpcId": VPC_ID, + } + mock_conn.get_private_graph_endpoint.return_value = {"vpcEndpointId": ENDPOINT_ID} + + operator = NeptuneCreatePrivateGraphEndpointOperator( + task_id="test_task", + graph_identifier=GRAPH_ID, + vpc_id=VPC_ID, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneDeletePrivateGraphEndpointOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -548,6 +601,31 @@ def test_execute_complete_success(self): # Verify the method completes without error and logs the endpoint_id + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn): + mock_conn.delete_private_graph_endpoint.return_value = { + "status": "DELETING", + "vpcEndpointId": ENDPOINT_ID, + "vpcId": VPC_ID, + } + + operator = NeptuneDeletePrivateGraphEndpointOperator( + task_id="test_task", + graph_identifier=GRAPH_ID, + vpc_id=VPC_ID, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneDeleteGraphOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -704,6 +782,31 @@ def test_delete_graph_other_client_error(self, mock_conn): with pytest.raises(NeptuneGraphDeletionFailedError): operator.execute(None) + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn): + mock_conn.delete_graph.return_value = { + "id": GRAPH_ID, + "name": GRAPH_NAME, + "status": "DELETING", + } + + operator = NeptuneDeleteGraphOperator( + task_id="test_task", + graph_id=GRAPH_ID, + skip_snapshot=True, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneCreateGraphWithImportOperator: IMPORT_TASK_ID = "import-task-12345" @@ -971,6 +1074,57 @@ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn): assert exc_info.value.method_name == "defer_wait_for_task" assert exc_info.value.kwargs == {"import_task_id": self.IMPORT_TASK_ID} + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn): + mock_conn.create_graph_using_import_task.return_value = { + "graphId": GRAPH_ID, + "taskId": self.IMPORT_TASK_ID, + "status": "IMPORTING", + } + + operator = NeptuneCreateGraphWithImportOperator( + task_id="test_task", + graph_name=GRAPH_NAME, + vector_search_config={"dimension": 128}, + source=SOURCE_S3_URI, + role_arn=ROLE_ARN, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_defer_wait_for_task_trigger_receives_hook_configuration(self, mock_conn): + operator = NeptuneCreateGraphWithImportOperator( + task_id="test_task", + graph_name=GRAPH_NAME, + vector_search_config={"dimension": 128}, + source=SOURCE_S3_URI, + role_arn=ROLE_ARN, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.defer_wait_for_task( + import_task_id=self.IMPORT_TASK_ID, + context=None, + event={"status": "success"}, + ) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + TASK_ID = "import-task-id-12345" @@ -1184,6 +1338,33 @@ def test_execute_complete_success(self): assert result == {"graph_id": GRAPH_ID, "import_task_id": TASK_ID} + @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneImportTaskLink.persist") + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_persist): + mock_conn.start_import_task.return_value = { + "taskId": TASK_ID, + "graphId": GRAPH_ID, + "status": "IMPORTING", + } + + operator = NeptuneStartImportTaskOperator( + task_id="test_task", + graph_identifier=GRAPH_ID, + role_arn=ROLE_ARN, + source=SOURCE_S3_URI, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestNeptuneCancelImportTaskOperator: @mock.patch.object(NeptuneAnalyticsHook, "conn") @@ -1275,3 +1456,27 @@ def test_execute_complete_success(self): result = operator.execute_complete(None, event) assert result == {"import_task_id": TASK_ID} + + @mock.patch.object(NeptuneAnalyticsHook, "conn") + def test_deferred_trigger_receives_hook_configuration(self, mock_conn): + mock_conn.cancel_import_task.return_value = { + "taskId": TASK_ID, + "graphId": GRAPH_ID, + "status": "CANCELLING", + } + + operator = NeptuneCancelImportTaskOperator( + task_id="test_task", + import_task_id=TASK_ID, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py index c6b114057eba9..000cc710a887f 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py @@ -22,7 +22,7 @@ from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook from airflow.providers.amazon.aws.sensors.mwaa import MwaaDagRunSensor, MwaaTaskSensor -from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred from airflow.utils.state import DagRunState, TaskInstanceState SENSOR_DAG_RUN_KWARGS = { @@ -46,6 +46,10 @@ "max_retries": 100, } +REGION_NAME = "eu-west-2" +VERIFY = False +BOTOCORE_CONFIG = {"read_timeout": 42} + SENSOR_STATE_KWARGS = { "success_states": ["a", "b"], "failure_states": ["c", "d"], @@ -109,6 +113,22 @@ def test_execute_complete_success(self): success_event = {"status": "success", "dag_run_id": "test_run"} sensor.execute_complete({}, success_event) # should not raise + def test_deferred_trigger_receives_hook_configuration(self): + sensor = MwaaDagRunSensor( + **{**SENSOR_DAG_RUN_KWARGS, "deferrable": True}, + **SENSOR_STATE_KWARGS, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + sensor.execute({}) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + class TestMwaaTaskSuccessSensor: def test_init_success(self): @@ -159,3 +179,19 @@ def test_execute_complete_success(self): sensor = MwaaTaskSensor(**SENSOR_TASK_KWARGS, **SENSOR_STATE_KWARGS) success_event = {"status": "success", "task_id": "test_task"} sensor.execute_complete({}, success_event) # should not raise + + def test_deferred_trigger_receives_hook_configuration(self): + sensor = MwaaTaskSensor( + **SENSOR_TASK_KWARGS, + **SENSOR_STATE_KWARGS, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + sensor.execute({}) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py index 4b714d1e90dda..9b4bdf80072c1 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py @@ -23,8 +23,12 @@ from airflow.providers.amazon.aws.hooks.ssm import SsmHook from airflow.providers.amazon.aws.sensors.ssm import SsmRunCommandCompletedSensor +from airflow.providers.common.compat.sdk import TaskDeferred COMMAND_ID = "123e4567-e89b-12d3-a456-426614174000" +REGION_NAME = "eu-west-2" +VERIFY = False +BOTOCORE_CONFIG = {"read_timeout": 42} @pytest.fixture @@ -144,3 +148,19 @@ def test_sensor_passes_parameter_to_trigger(self, mock_trigger_class): assert call_kwargs["command_id"] == COMMAND_ID assert call_kwargs["fail_on_nonzero_exit"] is False + + def test_deferred_trigger_receives_hook_configuration(self): + sensor = self.SENSOR( + **self.default_op_kwarg, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + + with pytest.raises(TaskDeferred) as exc_info: + sensor.execute({}) + + assert exc_info.value.trigger.region_name == REGION_NAME + assert exc_info.value.trigger.verify == VERIFY + assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG From fc53f763a30cc6a3629c1b80297dcab5e54dbec9 Mon Sep 17 00:00:00 2001 From: SEPURI-SAI-KRISHNA Date: Thu, 27 Aug 2026 14:24:14 +0530 Subject: [PATCH 2/2] Document the verify parameter for Neptune Analytics operators Every Neptune Analytics operator accepts `verify` through the shared AWS base class, but none of the seven class docstrings mentioned it, so the rendered provider docs gave users no way to discover it. Two of those docstrings even carried a stray blank line where the entry belonged. The deferral tests now compare the trigger's serialized payload rather than its attributes. Serialization is what actually crosses into the triggerer process, and it passes values through `prune_dict`, so an attribute-level assertion can pass while the setting is silently dropped on the way there. This matches the assertion style already used for the Neptune cluster operators. --- .../amazon/aws/operators/neptune_analytics.py | 16 ++- .../unit/amazon/aws/operators/test_mwaa.py | 16 ++- .../aws/operators/test_neptune_analytics.py | 99 ++++++++++++++----- .../unit/amazon/aws/sensors/test_mwaa.py | 33 +++++-- .../tests/unit/amazon/aws/sensors/test_ssm.py | 13 ++- 5 files changed, 139 insertions(+), 38 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py index 239906e5e9358..f8cc41881bd48 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py @@ -82,6 +82,8 @@ class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]): empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. See: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html :return: dictionary with Neptune graph id @@ -233,6 +235,8 @@ class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalytics empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. See: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html :return: dictionary with Neptune graph id @@ -379,7 +383,8 @@ class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalytics empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. - + :param verify: Whether or not to verify SSL certificates. See: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html :return: dictionary with Neptune graph id @@ -482,7 +487,8 @@ class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]): empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. - + :param verify: Whether or not to verify SSL certificates. See: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html :return: dictionary with Neptune graph id @@ -594,6 +600,8 @@ class NeptuneCreateGraphWithImportOperator(AwsBaseOperator[NeptuneAnalyticsHook] empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. See: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html :return: dictionary with Neptune graph id @@ -840,6 +848,8 @@ class NeptuneStartImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]): empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. See: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html :return: dictionary with Neptune graph id @@ -983,6 +993,8 @@ class NeptuneCancelImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]): empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. See: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html :return: dictionary with Neptune graph id diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py index e2232f77384fd..4e91ab1e1b472 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py @@ -133,6 +133,16 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_hook): with pytest.raises(TaskDeferred) as exc_info: op.execute({}) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": OP_KWARGS["waiter_delay"], + "waiter_max_attempts": OP_KWARGS["waiter_max_attempts"], + "aws_conn_id": "aws_default", + "external_env_name": OP_KWARGS["env_name"], + "external_dag_id": OP_KWARGS["trigger_dag_id"], + "external_dag_run_id": HOOK_RETURN_VALUE["RestApiResponse"]["dag_run_id"], + "success_states": None, + "failure_states": None, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py b/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py index f5aea2c2fc95a..ba92749660882 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py @@ -241,9 +241,15 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_pers with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } class TestNeptuneCreatePrivateGraphEndpointOperator: @@ -466,9 +472,16 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_pers with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "vpc_id": VPC_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } class TestNeptuneDeletePrivateGraphEndpointOperator: @@ -622,9 +635,17 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn): with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "vpc_id": VPC_ID, + "endpoint_id": ENDPOINT_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } class TestNeptuneDeleteGraphOperator: @@ -803,9 +824,15 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn): with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } class TestNeptuneCreateGraphWithImportOperator: @@ -1097,9 +1124,15 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn): with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "graph_id": GRAPH_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } @mock.patch.object(NeptuneAnalyticsHook, "conn") def test_defer_wait_for_task_trigger_receives_hook_configuration(self, mock_conn): @@ -1121,9 +1154,15 @@ def test_defer_wait_for_task_trigger_receives_hook_configuration(self, mock_conn event={"status": "success"}, ) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "import_task_id": self.IMPORT_TASK_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } TASK_ID = "import-task-id-12345" @@ -1361,9 +1400,15 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn, mock_pers with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "import_task_id": TASK_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } class TestNeptuneCancelImportTaskOperator: @@ -1477,6 +1522,12 @@ def test_deferred_trigger_receives_hook_configuration(self, mock_conn): with pytest.raises(TaskDeferred) as exc_info: operator.execute(None) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": 30, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + "task_identifier": TASK_ID, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py index 000cc710a887f..d1da019247ebc 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py @@ -125,9 +125,19 @@ def test_deferred_trigger_receives_hook_configuration(self): with pytest.raises(TaskDeferred) as exc_info: sensor.execute({}) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": SENSOR_DAG_RUN_KWARGS["poke_interval"], + "waiter_max_attempts": SENSOR_DAG_RUN_KWARGS["max_retries"], + "aws_conn_id": "aws_default", + "external_env_name": SENSOR_DAG_RUN_KWARGS["external_env_name"], + "external_dag_id": SENSOR_DAG_RUN_KWARGS["external_dag_id"], + "external_dag_run_id": SENSOR_DAG_RUN_KWARGS["external_dag_run_id"], + "success_states": set(SENSOR_STATE_KWARGS["success_states"]), + "failure_states": set(SENSOR_STATE_KWARGS["failure_states"]), + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } class TestMwaaTaskSuccessSensor: @@ -192,6 +202,17 @@ def test_deferred_trigger_receives_hook_configuration(self): with pytest.raises(TaskDeferred) as exc_info: sensor.execute({}) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": SENSOR_TASK_KWARGS["poke_interval"], + "waiter_max_attempts": SENSOR_TASK_KWARGS["max_retries"], + "aws_conn_id": "aws_default", + "external_env_name": SENSOR_TASK_KWARGS["external_env_name"], + "external_dag_id": SENSOR_TASK_KWARGS["external_dag_id"], + "external_dag_run_id": SENSOR_TASK_KWARGS["external_dag_run_id"], + "external_task_id": SENSOR_TASK_KWARGS["external_task_id"], + "success_states": set(SENSOR_STATE_KWARGS["success_states"]), + "failure_states": set(SENSOR_STATE_KWARGS["failure_states"]), + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + } diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py index 9b4bdf80072c1..a8472776b4cbe 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_ssm.py @@ -161,6 +161,13 @@ def test_deferred_trigger_receives_hook_configuration(self): with pytest.raises(TaskDeferred) as exc_info: sensor.execute({}) - assert exc_info.value.trigger.region_name == REGION_NAME - assert exc_info.value.trigger.verify == VERIFY - assert exc_info.value.trigger.botocore_config == BOTOCORE_CONFIG + assert exc_info.value.trigger.serialize()[1] == { + "waiter_delay": self.default_op_kwarg["poke_interval"], + "waiter_max_attempts": self.default_op_kwarg["max_retries"], + "aws_conn_id": "aws_default", + "command_id": COMMAND_ID, + "fail_on_nonzero_exit": True, + "region_name": REGION_NAME, + "verify": VERIFY, + "botocore_config": BOTOCORE_CONFIG, + }