From f6e61c1f73fea8158e9ccc59c46dc87deb1df7d6 Mon Sep 17 00:00:00 2001 From: heitorlessa Date: Wed, 1 Mar 2023 21:06:02 +0100 Subject: [PATCH 1/3] chore: add tests for static_pk_value feature --- tests/functional/idempotency/conftest.py | 62 +++++++++++++++++++ .../idempotency/test_idempotency.py | 31 ++++++++++ 2 files changed, 93 insertions(+) diff --git a/tests/functional/idempotency/conftest.py b/tests/functional/idempotency/conftest.py index 7e5fa0e7c61..483a7b60d4b 100644 --- a/tests/functional/idempotency/conftest.py +++ b/tests/functional/idempotency/conftest.py @@ -215,6 +215,13 @@ def persistence_store_compound(config): return DynamoDBPersistenceLayer(table_name=TABLE_NAME, boto_config=config, key_attr="id", sort_key_attr="sk") +@pytest.fixture +def persistence_store_compound_static_pk_value(config, static_pk_value): + return DynamoDBPersistenceLayer( + table_name=TABLE_NAME, boto_config=config, key_attr="id", sort_key_attr="sk", static_pk_value=static_pk_value + ) + + @pytest.fixture def idempotency_config(config, request, default_jmespath): return IdempotencyConfig( @@ -246,3 +253,58 @@ def _func_echo_decoder(self, value): @pytest.fixture def mock_function(): return mock.MagicMock() + + +@pytest.fixture +def static_pk_value(): + return "static-value" + + +@pytest.fixture +def expected_params_update_item_compound_key_static_pk_value( + serialized_lambda_response, hashed_idempotency_key, static_pk_value +): + return { + "ExpressionAttributeNames": { + "#expiry": "expiration", + "#response_data": "data", + "#status": "status", + }, + "ExpressionAttributeValues": { + ":expiry": {"N": stub.ANY}, + ":response_data": {"S": serialized_lambda_response}, + ":status": {"S": "COMPLETED"}, + }, + "Key": {"id": {"S": static_pk_value}, "sk": {"S": hashed_idempotency_key}}, + "TableName": "TEST_TABLE", + "UpdateExpression": "SET #response_data = :response_data, " "#expiry = :expiry, #status = :status", + } + + +@pytest.fixture +def expected_params_put_item_compound_key_static_pk_value(hashed_idempotency_key, static_pk_value): + return { + "ConditionExpression": ( + "attribute_not_exists(#id) OR #expiry < :now OR " + "(#status = :inprogress AND attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now_in_millis)" + ), + "ExpressionAttributeNames": { + "#id": "id", + "#expiry": "expiration", + "#status": "status", + "#in_progress_expiry": "in_progress_expiration", + }, + "ExpressionAttributeValues": { + ":now": {"N": stub.ANY}, + ":now_in_millis": {"N": stub.ANY}, + ":inprogress": {"S": "INPROGRESS"}, + }, + "Item": { + "expiration": {"N": stub.ANY}, + "in_progress_expiration": {"N": stub.ANY}, + "id": {"S": static_pk_value}, + "sk": {"S": hashed_idempotency_key}, + "status": {"S": "INPROGRESS"}, + }, + "TableName": "TEST_TABLE", + } diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py index dfc6b03b60c..68aeabeb50a 100644 --- a/tests/functional/idempotency/test_idempotency.py +++ b/tests/functional/idempotency/test_idempotency.py @@ -1504,3 +1504,34 @@ def lambda_handler(event, context): stubber.assert_no_pending_responses() stubber.deactivate() + + +@pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}], indirect=True) +def test_idempotent_lambda_compound_static_pk_value_has_correct_pk( + idempotency_config: IdempotencyConfig, + persistence_store_compound_static_pk_value: DynamoDBPersistenceLayer, + lambda_apigw_event, + expected_params_put_item_compound_key_static_pk_value, + expected_params_update_item_compound_key_static_pk_value, + lambda_response, + lambda_context, +): + """ + Test idempotent decorator having a DynamoDBPersistenceLayer with a compound key and a static PK value + """ + + stubber = stub.Stubber(persistence_store_compound_static_pk_value._client) + ddb_response = {} + + stubber.add_response("put_item", ddb_response, expected_params_put_item_compound_key_static_pk_value) + stubber.add_response("update_item", ddb_response, expected_params_update_item_compound_key_static_pk_value) + stubber.activate() + + @idempotent(config=idempotency_config, persistence_store=persistence_store_compound_static_pk_value) + def lambda_handler(event, context): + return lambda_response + + lambda_handler(lambda_apigw_event, lambda_context) + + stubber.assert_no_pending_responses() + stubber.deactivate() From 10f8833a6ba26060a0f4ebbeafea874136b654b1 Mon Sep 17 00:00:00 2001 From: heitorlessa Date: Wed, 1 Mar 2023 21:12:27 +0100 Subject: [PATCH 2/3] fix: remove dict mutation for primary key causing regression --- .../idempotency/persistence/dynamodb.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py index b05d8216b50..ce3aac2425e 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py +++ b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py @@ -104,6 +104,21 @@ def __init__( super(DynamoDBPersistenceLayer, self).__init__() def _get_key(self, idempotency_key: str) -> dict: + """Build primary key attribute simple or composite based on params. + + When sort_key_attr is set, we must return a composite key with static_pk_value, + otherwise we use the idempotency key given. + + Parameters + ---------- + idempotency_key : str + idempotency key to use for simple primary key + + Returns + ------- + dict + simple or composite key for DynamoDB primary key + """ if self.sort_key_attr: return {self.key_attr: {"S": self.static_pk_value}, self.sort_key_attr: {"S": idempotency_key}} return {self.key_attr: {"S": idempotency_key}} @@ -145,8 +160,8 @@ def _get_record(self, idempotency_key) -> DataRecord: def _put_record(self, data_record: DataRecord) -> None: item = { + # get simple or composite primary key **self._get_key(data_record.idempotency_key), - self.key_attr: {"S": data_record.idempotency_key}, self.expiry_attr: {"N": str(data_record.expiry_timestamp)}, self.status_attr: {"S": data_record.status}, } From f87778a57e8a441a0c18f1ae6bcc7416ee53f9a7 Mon Sep 17 00:00:00 2001 From: heitorlessa Date: Wed, 1 Mar 2023 21:26:45 +0100 Subject: [PATCH 3/3] chore: reduce boilerplate --- tests/functional/idempotency/conftest.py | 37 +++++------------------- 1 file changed, 8 insertions(+), 29 deletions(-) diff --git a/tests/functional/idempotency/conftest.py b/tests/functional/idempotency/conftest.py index 483a7b60d4b..75834f76688 100644 --- a/tests/functional/idempotency/conftest.py +++ b/tests/functional/idempotency/conftest.py @@ -262,43 +262,22 @@ def static_pk_value(): @pytest.fixture def expected_params_update_item_compound_key_static_pk_value( - serialized_lambda_response, hashed_idempotency_key, static_pk_value + expected_params_update_item, hashed_idempotency_key, static_pk_value ): return { - "ExpressionAttributeNames": { - "#expiry": "expiration", - "#response_data": "data", - "#status": "status", - }, - "ExpressionAttributeValues": { - ":expiry": {"N": stub.ANY}, - ":response_data": {"S": serialized_lambda_response}, - ":status": {"S": "COMPLETED"}, - }, + # same as in any update_item transaction except the `Key` due to composite key value + **expected_params_update_item, "Key": {"id": {"S": static_pk_value}, "sk": {"S": hashed_idempotency_key}}, - "TableName": "TEST_TABLE", - "UpdateExpression": "SET #response_data = :response_data, " "#expiry = :expiry, #status = :status", } @pytest.fixture -def expected_params_put_item_compound_key_static_pk_value(hashed_idempotency_key, static_pk_value): +def expected_params_put_item_compound_key_static_pk_value( + expected_params_put_item, hashed_idempotency_key, static_pk_value +): return { - "ConditionExpression": ( - "attribute_not_exists(#id) OR #expiry < :now OR " - "(#status = :inprogress AND attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now_in_millis)" - ), - "ExpressionAttributeNames": { - "#id": "id", - "#expiry": "expiration", - "#status": "status", - "#in_progress_expiry": "in_progress_expiration", - }, - "ExpressionAttributeValues": { - ":now": {"N": stub.ANY}, - ":now_in_millis": {"N": stub.ANY}, - ":inprogress": {"S": "INPROGRESS"}, - }, + # same as in any put_item transaction except the `Item` due to composite key value + **expected_params_put_item, "Item": { "expiration": {"N": stub.ANY}, "in_progress_expiration": {"N": stub.ANY},