This is an automated email from the ASF dual-hosted git repository.
o-nikolas 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 32168fdf5ef Add deferrable mode to DynamoDBValueSensor (#69799)
32168fdf5ef is described below
commit 32168fdf5ef88c7f0efd0e1ba0c40f426869753e
Author: Harsh Gupta <[email protected]>
AuthorDate: Tue Jul 28 05:28:36 2026 +0530
Add deferrable mode to DynamoDBValueSensor (#69799)
DynamoDBValueSensor previously only supported poke/reschedule modes,
occupying a worker slot while waiting for a DynamoDB attribute value.
This adds the standard deferrable path used across the Amazon provider:
---
providers/amazon/docs/operators/dynamodb.rst | 1 +
providers/amazon/provider.yaml | 3 +
.../providers/amazon/aws/sensors/dynamodb.py | 39 +++++
.../providers/amazon/aws/triggers/dynamodb.py | 168 ++++++++++++++++++++
.../airflow/providers/amazon/get_provider_info.py | 4 +
.../tests/unit/amazon/aws/sensors/test_dynamodb.py | 54 +++++++
.../unit/amazon/aws/triggers/test_dynamodb.py | 175 +++++++++++++++++++++
7 files changed, 444 insertions(+)
diff --git a/providers/amazon/docs/operators/dynamodb.rst
b/providers/amazon/docs/operators/dynamodb.rst
index f802fcdb423..633b311fcf0 100644
--- a/providers/amazon/docs/operators/dynamodb.rst
+++ b/providers/amazon/docs/operators/dynamodb.rst
@@ -45,6 +45,7 @@ Wait on Amazon DynamoDB item attribute value match
Use the
:class:`~airflow.providers.amazon.aws.sensors.dynamodb.DynamoDBValueSensor`
to wait for the presence of a matching DynamoDB item's attribute/value pair.
+This sensor can also be run in deferrable mode by setting ``deferrable`` param
to ``True``.
Wait for a Single Attribute Value Match:
----------------------------------------
diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml
index 0121c5b5ee2..343398db658 100644
--- a/providers/amazon/provider.yaml
+++ b/providers/amazon/provider.yaml
@@ -827,6 +827,9 @@ triggers:
- integration-name: Amazon Comprehend
python-modules:
- airflow.providers.amazon.aws.triggers.comprehend
+ - integration-name: Amazon DynamoDB
+ python-modules:
+ - airflow.providers.amazon.aws.triggers.dynamodb
- integration-name: Amazon EC2
python-modules:
- airflow.providers.amazon.aws.triggers.ec2
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/sensors/dynamodb.py
b/providers/amazon/src/airflow/providers/amazon/aws/sensors/dynamodb.py
index 8bf1a67f56c..a3efa77b04f 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/dynamodb.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/dynamodb.py
@@ -17,13 +17,17 @@
from __future__ import annotations
from collections.abc import Iterable, Sequence
+from datetime import timedelta
from typing import TYPE_CHECKING, Any
from botocore.exceptions import ClientError
from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.triggers.dynamodb import
DynamoDBValueSensorTrigger
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
if TYPE_CHECKING:
from airflow.sdk import Context
@@ -44,6 +48,9 @@ class DynamoDBValueSensor(AwsBaseSensor[DynamoDBHook]):
:param attribute_value: DynamoDB attribute value
:param sort_key_name: (optional) DynamoDB sort key name
:param sort_key_value: (optional) DynamoDB sort key value
+ :param deferrable: If True, the sensor will operate in deferrable mode.
This mode requires aiobotocore
+ module to be installed.
+ (default: False, but can be overridden in config file by setting
default_deferrable to True)
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used.
If
running Airflow in a distributed manner and aws_conn_id is None or
@@ -76,6 +83,7 @@ class DynamoDBValueSensor(AwsBaseSensor[DynamoDBHook]):
attribute_value: str | Iterable[str],
sort_key_name: str | None = None,
sort_key_value: str | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
**kwargs: Any,
):
super().__init__(**kwargs)
@@ -86,6 +94,37 @@ class DynamoDBValueSensor(AwsBaseSensor[DynamoDBHook]):
self.attribute_value = attribute_value
self.sort_key_name = sort_key_name
self.sort_key_value = sort_key_value
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> Any:
+ if self.deferrable:
+ self.defer(
+ trigger=DynamoDBValueSensorTrigger(
+ table_name=self.table_name,
+ partition_key_name=self.partition_key_name,
+ partition_key_value=self.partition_key_value,
+ attribute_name=self.attribute_name,
+ attribute_value=self.attribute_value,
+ sort_key_name=self.sort_key_name,
+ sort_key_value=self.sort_key_value,
+ waiter_delay=int(self.poke_interval),
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ botocore_config=self.botocore_config,
+ ),
+ method_name="execute_complete",
+ timeout=timedelta(seconds=self.timeout),
+ )
+ else:
+ super().execute(context=context)
+
+ def execute_complete(self, context: Context, event: dict | None = None) ->
None:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event["status"] != "success":
+ raise RuntimeError(f"Trigger error: event is {validated_event}")
+ self.log.info("DynamoDB attribute value match found; sensor complete.")
def poke(self, context: Context) -> bool:
"""Test DynamoDB item for matching attribute value."""
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/triggers/dynamodb.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/dynamodb.py
new file mode 100644
index 00000000000..f1008b79f4d
--- /dev/null
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/dynamodb.py
@@ -0,0 +1,168 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator, Iterable
+from functools import cached_property
+from typing import Any
+
+from boto3.dynamodb.types import TypeDeserializer
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+else:
+ from airflow.triggers.base import ( # type: ignore
+ BaseTrigger as BaseEventTrigger,
+ TriggerEvent,
+ )
+
+
+class DynamoDBValueSensorTrigger(BaseEventTrigger):
+ """
+ Asynchronously poll a DynamoDB item until the given attribute matches one
of the expected values.
+
+ The polling uses the low-level DynamoDB client (aiobotocore exposes
clients only, not the
+ boto3 resource API used by ``DynamoDBHook``), so keys are sent as typed
string attributes,
+ matching the string-typed key parameters of
:class:`~airflow.providers.amazon.aws.sensors.dynamodb.DynamoDBValueSensor`.
+
+ :param table_name: DynamoDB table name
+ :param partition_key_name: DynamoDB partition key name
+ :param partition_key_value: DynamoDB partition key value
+ :param attribute_name: DynamoDB attribute name
+ :param attribute_value: expected DynamoDB attribute value (or values, any
of which matches)
+ :param sort_key_name: (optional) DynamoDB sort key name
+ :param sort_key_value: (optional) DynamoDB sort key value
+ :param waiter_delay: The time in seconds to wait between DynamoDB API calls
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ :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
+ """
+
+ def __init__(
+ self,
+ table_name: str,
+ partition_key_name: str,
+ partition_key_value: str,
+ attribute_name: str,
+ attribute_value: str | Iterable[str],
+ sort_key_name: str | None = None,
+ sort_key_value: str | None = None,
+ waiter_delay: int = 60,
+ aws_conn_id: str | None = "aws_default",
+ region_name: str | None = None,
+ verify: bool | str | None = None,
+ botocore_config: dict | None = None,
+ ):
+ self.table_name = table_name
+ self.partition_key_name = partition_key_name
+ self.partition_key_value = partition_key_value
+ self.attribute_name = attribute_name
+ self.attribute_value = (
+ [attribute_value] if isinstance(attribute_value, str) else
list(attribute_value)
+ )
+ self.sort_key_name = sort_key_name
+ self.sort_key_value = sort_key_value
+ self.waiter_delay = waiter_delay
+ self.aws_conn_id = aws_conn_id
+ self.region_name = region_name
+ self.verify = verify
+ self.botocore_config = botocore_config
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+ self.__class__.__module__ + "." + self.__class__.__qualname__,
+ {
+ "table_name": self.table_name,
+ "partition_key_name": self.partition_key_name,
+ "partition_key_value": self.partition_key_value,
+ "attribute_name": self.attribute_name,
+ "attribute_value": self.attribute_value,
+ "sort_key_name": self.sort_key_name,
+ "sort_key_value": self.sort_key_value,
+ "waiter_delay": self.waiter_delay,
+ "aws_conn_id": self.aws_conn_id,
+ "region_name": self.region_name,
+ "verify": self.verify,
+ "botocore_config": self.botocore_config,
+ },
+ )
+
+ @cached_property
+ def hook(self) -> AwsBaseHook:
+ # DynamoDBHook is resource-based, but async connections are only
available for
+ # clients, so the trigger talks to DynamoDB through a client-type hook
instead.
+ return AwsBaseHook(
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ config=self.botocore_config,
+ client_type="dynamodb",
+ )
+
+ @property
+ def key(self) -> dict[str, Any]:
+ key = {self.partition_key_name: {"S": self.partition_key_value}}
+ if self.sort_key_name and self.sort_key_value:
+ key[self.sort_key_name] = {"S": self.sort_key_value}
+ return key
+
+ async def poke(self, client: Any) -> bool:
+ """Test the DynamoDB item for a matching attribute value, mirroring
the sensor's poke."""
+ self.log.info(
+ "Checking table %s for item with key %s, waiting for attribute %s
to be one of %s",
+ self.table_name,
+ self.key,
+ self.attribute_name,
+ self.attribute_value,
+ )
+ try:
+ response = await client.get_item(TableName=self.table_name,
Key=self.key)
+ except ClientError as err:
+ # Same tolerance as the sensor's poke: log and keep trying until
the task times out.
+ self.log.error(
+ "Couldn't get %s from table %s.\nError Code: %s\nError
Message: %s",
+ self.key,
+ self.table_name,
+ err.response["Error"]["Code"],
+ err.response["Error"]["Message"],
+ )
+ return False
+
+ try:
+ typed_attribute_value = response["Item"][self.attribute_name]
+ except KeyError:
+ return False
+ item_attribute_value =
TypeDeserializer().deserialize(typed_attribute_value)
+ self.log.info("Got: %s = %s", self.attribute_name,
item_attribute_value)
+ return item_attribute_value in self.attribute_value
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ while True:
+ # This loop runs until the timeout set in the sensor's self.defer
call is reached.
+ async with await self.hook.get_async_conn() as client:
+ if await self.poke(client=client):
+ yield TriggerEvent({"status": "success"})
+ return
+ await asyncio.sleep(self.waiter_delay)
diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
index 5314709a5df..a3e609aa5a8 100644
--- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
+++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
@@ -926,6 +926,10 @@ def get_provider_info():
"integration-name": "Amazon Comprehend",
"python-modules":
["airflow.providers.amazon.aws.triggers.comprehend"],
},
+ {
+ "integration-name": "Amazon DynamoDB",
+ "python-modules":
["airflow.providers.amazon.aws.triggers.dynamodb"],
+ },
{
"integration-name": "Amazon EC2",
"python-modules":
["airflow.providers.amazon.aws.triggers.ec2"],
diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_dynamodb.py
b/providers/amazon/tests/unit/amazon/aws/sensors/test_dynamodb.py
index e8ea324ff8c..46999c4c81a 100644
--- a/providers/amazon/tests/unit/amazon/aws/sensors/test_dynamodb.py
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_dynamodb.py
@@ -19,10 +19,13 @@ from __future__ import annotations
from unittest import mock
+import pytest
from moto import mock_aws
from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
from airflow.providers.amazon.aws.sensors.dynamodb import DynamoDBValueSensor
+from airflow.providers.amazon.aws.triggers.dynamodb import
DynamoDBValueSensorTrigger
+from airflow.providers.common.compat.sdk import TaskDeferred
class TestDynamoDBValueSensor:
@@ -120,6 +123,57 @@ class TestDynamoDBValueSensor:
self.sensor_pk.partition_key_name = "no such key"
assert self.sensor_pk.poke(None) is False
+ def test_sensor_deferrable(self):
+ sensor = DynamoDBValueSensor(
+ task_id="dynamodb_value_sensor_deferrable",
+ table_name=self.table_name,
+ partition_key_name=self.pk_name,
+ partition_key_value=self.pk_value,
+ attribute_name=self.attribute_name,
+ attribute_value=self.attribute_value,
+ sort_key_name=self.sk_name,
+ sort_key_value=self.sk_value,
+ deferrable=True,
+ )
+ with pytest.raises(TaskDeferred) as defer:
+ sensor.execute({})
+ trigger = defer.value.trigger
+ assert isinstance(trigger, DynamoDBValueSensorTrigger)
+ assert trigger.table_name == self.table_name
+ assert trigger.partition_key_name == self.pk_name
+ assert trigger.partition_key_value == self.pk_value
+ assert trigger.attribute_name == self.attribute_name
+ assert trigger.attribute_value == [self.attribute_value]
+ assert trigger.sort_key_name == self.sk_name
+ assert trigger.sort_key_value == self.sk_value
+
+ def test_execute_complete(self):
+ sensor = DynamoDBValueSensor(
+ task_id="dynamodb_value_sensor_deferrable",
+ table_name=self.table_name,
+ partition_key_name=self.pk_name,
+ partition_key_value=self.pk_value,
+ attribute_name=self.attribute_name,
+ attribute_value=self.attribute_value,
+ deferrable=True,
+ )
+ sensor.execute_complete(context={}, event={"status": "success"})
+
+ def test_fail_execute_complete(self):
+ sensor = DynamoDBValueSensor(
+ task_id="dynamodb_value_sensor_deferrable",
+ table_name=self.table_name,
+ partition_key_name=self.pk_name,
+ partition_key_value=self.pk_value,
+ attribute_name=self.attribute_name,
+ attribute_value=self.attribute_value,
+ deferrable=True,
+ )
+ event = {"status": "failed"}
+ message = f"Trigger error: event is {event}"
+ with pytest.raises(RuntimeError, match=message):
+ sensor.execute_complete(context={}, event=event)
+
class TestDynamoDBMultipleValuesSensor:
def setup_method(self):
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_dynamodb.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_dynamodb.py
new file mode 100644
index 00000000000..afcd1991384
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_dynamodb.py
@@ -0,0 +1,175 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+from unittest.mock import AsyncMock
+
+import pytest
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.triggers.dynamodb import
DynamoDBValueSensorTrigger
+from airflow.triggers.base import TriggerEvent
+
+TEST_TABLE_NAME = "test-table"
+TEST_PK_NAME = "PK"
+TEST_PK_VALUE = "PKTest"
+TEST_SK_NAME = "SK"
+TEST_SK_VALUE = "SKTest"
+TEST_ATTRIBUTE_NAME = "Foo"
+TEST_ATTRIBUTE_VALUE = "Bar"
+TEST_WAITER_DELAY = 1
+TEST_CONN_ID = "test-conn-id"
+TEST_REGION_NAME = "eu-central-1"
+
+
+def _build_trigger(**kwargs) -> DynamoDBValueSensorTrigger:
+ params = {
+ "table_name": TEST_TABLE_NAME,
+ "partition_key_name": TEST_PK_NAME,
+ "partition_key_value": TEST_PK_VALUE,
+ "attribute_name": TEST_ATTRIBUTE_NAME,
+ "attribute_value": TEST_ATTRIBUTE_VALUE,
+ "waiter_delay": TEST_WAITER_DELAY,
+ "aws_conn_id": TEST_CONN_ID,
+ "region_name": TEST_REGION_NAME,
+ **kwargs,
+ }
+ return DynamoDBValueSensorTrigger(**params)
+
+
+class TestDynamoDBValueSensorTrigger:
+ def test_serialize(self):
+ trigger = _build_trigger(sort_key_name=TEST_SK_NAME,
sort_key_value=TEST_SK_VALUE)
+
+ class_path, args = trigger.serialize()
+ assert class_path ==
"airflow.providers.amazon.aws.triggers.dynamodb.DynamoDBValueSensorTrigger"
+ assert args["table_name"] == TEST_TABLE_NAME
+ assert args["partition_key_name"] == TEST_PK_NAME
+ assert args["partition_key_value"] == TEST_PK_VALUE
+ assert args["attribute_name"] == TEST_ATTRIBUTE_NAME
+ # a single string value is normalized to a list so serialization
round-trips
+ assert args["attribute_value"] == [TEST_ATTRIBUTE_VALUE]
+ assert args["sort_key_name"] == TEST_SK_NAME
+ assert args["sort_key_value"] == TEST_SK_VALUE
+ assert args["waiter_delay"] == TEST_WAITER_DELAY
+ assert args["aws_conn_id"] == TEST_CONN_ID
+ assert args["region_name"] == TEST_REGION_NAME
+ assert args["verify"] is None
+ assert args["botocore_config"] is None
+
+ def test_serialize_generic_hook_params(self):
+ trigger = _build_trigger(verify=False,
botocore_config={"read_timeout": 99})
+ _, args = trigger.serialize()
+ assert args["verify"] is False
+ assert args["botocore_config"] == {"read_timeout": 99}
+
+ hook = trigger.hook
+ assert hook.client_type == "dynamodb"
+ assert hook.aws_conn_id == TEST_CONN_ID
+ assert hook._region_name == TEST_REGION_NAME
+ assert hook._verify is False
+ assert hook._config.read_timeout == 99
+
+ def test_key_with_and_without_sort_key(self):
+ trigger = _build_trigger()
+ assert trigger.key == {TEST_PK_NAME: {"S": TEST_PK_VALUE}}
+
+ trigger = _build_trigger(sort_key_name=TEST_SK_NAME,
sort_key_value=TEST_SK_VALUE)
+ assert trigger.key == {
+ TEST_PK_NAME: {"S": TEST_PK_VALUE},
+ TEST_SK_NAME: {"S": TEST_SK_VALUE},
+ }
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_async_conn")
+ async def test_run_success_when_value_matches(self, mock_async_conn):
+ client = AsyncMock()
+ mock_async_conn.return_value.__aenter__.return_value = client
+ client.get_item.return_value = {
+ "Item": {
+ TEST_PK_NAME: {"S": TEST_PK_VALUE},
+ TEST_ATTRIBUTE_NAME: {"S": TEST_ATTRIBUTE_VALUE},
+ }
+ }
+
+ trigger = _build_trigger()
+ generator = trigger.run()
+ response = await generator.asend(None)
+
+ client.get_item.assert_called_once_with(
+ TableName=TEST_TABLE_NAME, Key={TEST_PK_NAME: {"S": TEST_PK_VALUE}}
+ )
+ assert response == TriggerEvent({"status": "success"})
+
+ @pytest.mark.asyncio
+ @mock.patch("asyncio.sleep")
+
@mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_async_conn")
+ async def test_run_keeps_polling_until_value_matches(self,
mock_async_conn, mock_sleep):
+ mock_sleep.return_value = True
+ client = AsyncMock()
+ mock_async_conn.return_value.__aenter__.return_value = client
+ client.get_item.side_effect = [
+ {}, # item does not exist yet
+ {"Item": {TEST_PK_NAME: {"S": TEST_PK_VALUE}, TEST_ATTRIBUTE_NAME:
{"S": "wrong"}}},
+ {"Item": {TEST_PK_NAME: {"S": TEST_PK_VALUE}, TEST_ATTRIBUTE_NAME:
{"S": TEST_ATTRIBUTE_VALUE}}},
+ ]
+
+ trigger = _build_trigger()
+ generator = trigger.run()
+ response = await generator.asend(None)
+
+ assert client.get_item.call_count == 3
+ assert response == TriggerEvent({"status": "success"})
+
+ @pytest.mark.asyncio
+ @mock.patch("asyncio.sleep")
+
@mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_async_conn")
+ async def test_run_tolerates_client_error(self, mock_async_conn,
mock_sleep):
+ # Same tolerance as the sensor's poke: a ClientError means "not there
yet", not failure.
+ mock_sleep.return_value = True
+ client = AsyncMock()
+ mock_async_conn.return_value.__aenter__.return_value = client
+ client.get_item.side_effect = [
+ ClientError(
+ error_response={"Error": {"Code": "ResourceNotFoundException",
"Message": "no table"}},
+ operation_name="GetItem",
+ ),
+ {"Item": {TEST_PK_NAME: {"S": TEST_PK_VALUE}, TEST_ATTRIBUTE_NAME:
{"S": TEST_ATTRIBUTE_VALUE}}},
+ ]
+
+ trigger = _build_trigger()
+ generator = trigger.run()
+ response = await generator.asend(None)
+
+ assert client.get_item.call_count == 2
+ assert response == TriggerEvent({"status": "success"})
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_async_conn")
+ async def test_run_success_with_multiple_expected_values(self,
mock_async_conn):
+ client = AsyncMock()
+ mock_async_conn.return_value.__aenter__.return_value = client
+ client.get_item.return_value = {
+ "Item": {TEST_PK_NAME: {"S": TEST_PK_VALUE}, TEST_ATTRIBUTE_NAME:
{"S": "Bar2"}}
+ }
+
+ trigger = _build_trigger(attribute_value=["Bar1", "Bar2", "Bar3"])
+ generator = trigger.run()
+ response = await generator.asend(None)
+
+ assert response == TriggerEvent({"status": "success"})
