From 60c81dbe48635bf157ec9781612b8a2dc108108b Mon Sep 17 00:00:00 2001 From: Utkarsh Sharma Date: Fri, 14 Apr 2023 22:13:59 +0530 Subject: [PATCH] DynamoDBHook - waiter_path() to consider `resource_type` or `client_type` (#30595) * Add while initilizing * Add while initilizing * Add logic to pick either client_type or resource_type * Add test case * Assert expected path --- airflow/providers/amazon/aws/hooks/base_aws.py | 3 ++- tests/providers/amazon/aws/hooks/test_dynamodb.py | 7 +++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 3395990fc343e..77e372e88bbc0 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -794,7 +794,8 @@ def test_connection(self): @cached_property def waiter_path(self) -> PathLike[str] | None: - path = Path(__file__).parents[1].joinpath(f"waiters/{self.client_type}.json").resolve() + filename = self.client_type if self.client_type else self.resource_type + path = Path(__file__).parents[1].joinpath(f"waiters/{filename}.json").resolve() return path if path.exists() else None def get_waiter(self, waiter_name: str, parameters: dict[str, str] | None = None) -> Waiter: diff --git a/tests/providers/amazon/aws/hooks/test_dynamodb.py b/tests/providers/amazon/aws/hooks/test_dynamodb.py index 8c5886639c0cc..15fac620f1c35 100644 --- a/tests/providers/amazon/aws/hooks/test_dynamodb.py +++ b/tests/providers/amazon/aws/hooks/test_dynamodb.py @@ -18,6 +18,7 @@ from __future__ import annotations import uuid +from unittest import mock from moto import mock_dynamodb @@ -55,3 +56,9 @@ def test_insert_batch_items_dynamodb_table(self): table.meta.client.get_waiter("table_exists").wait(TableName="test_airflow") assert table.item_count == 10 + + @mock.patch("pathlib.Path.exists", return_value=True) + def test_waiter_path_generated_from_resource_type(self, _): + hook = DynamoDBHook(aws_conn_id="aws_default") + path = hook.waiter_path + assert path.as_uri().endswith("/airflow/airflow/providers/amazon/aws/waiters/dynamodb.json")