diff --git a/aws_lambda_powertools/utilities/idempotency/base.py b/aws_lambda_powertools/utilities/idempotency/base.py index 9b54421e40b..6978cd778de 100644 --- a/aws_lambda_powertools/utilities/idempotency/base.py +++ b/aws_lambda_powertools/utilities/idempotency/base.py @@ -239,18 +239,15 @@ def _handle_for_status(self, data_record: DataRecord) -> Any | None: f"Execution already in progress with idempotency key: " f"{self.persistence_store.event_key_jmespath}={data_record.idempotency_key}", ) - response_dict: dict | None = data_record.response_json_as_dict() - if response_dict is not None: - serialized_response = self.output_serializer.from_dict(response_dict) - if self.config.response_hook is not None: - logger.debug("Response hook configured, invoking function") - return self.config.response_hook( - serialized_response, - data_record, - ) - return serialized_response - return None + response_dict = data_record.response_json_as_dict() + serialized_response = self.output_serializer.from_dict(response_dict) if response_dict else None + + if self.config.response_hook: + logger.debug("Response hook configured, invoking function") + return self.config.response_hook(serialized_response, data_record) + + return serialized_response def _get_function_response(self): try: diff --git a/tests/e2e/idempotency/handlers/response_hook.py b/tests/e2e/idempotency/handlers/response_hook.py new file mode 100644 index 00000000000..b56d09f25b9 --- /dev/null +++ b/tests/e2e/idempotency/handlers/response_hook.py @@ -0,0 +1,29 @@ +import os + +from aws_lambda_powertools.utilities.idempotency import ( + DynamoDBPersistenceLayer, + IdempotencyConfig, + idempotent, +) +from aws_lambda_powertools.utilities.idempotency.persistence.datarecord import ( + DataRecord, +) + +TABLE_NAME = os.getenv("IdempotencyTable", "") +persistence_layer = DynamoDBPersistenceLayer(table_name=TABLE_NAME) + + +def my_response_hook(response: dict, idempotent_data: DataRecord) -> dict: + # Return inserted Header data into the Idempotent Response + response["x-response-hook"] = idempotent_data.idempotency_key + + # Must return the response here + return response + + +config = IdempotencyConfig(response_hook=my_response_hook) + + +@idempotent(config=config, persistence_store=persistence_layer) +def lambda_handler(event, context): + return {"message": "first_response"} diff --git a/tests/e2e/idempotency/infrastructure.py b/tests/e2e/idempotency/infrastructure.py index bcc35005549..6ca14d576ef 100644 --- a/tests/e2e/idempotency/infrastructure.py +++ b/tests/e2e/idempotency/infrastructure.py @@ -20,6 +20,7 @@ def create_resources(self): table.grant_read_write_data(functions["FunctionThreadSafetyHandler"]) table.grant_read_write_data(functions["OptionalIdempotencyKeyHandler"]) table.grant_read_write_data(functions["PayloadTamperingValidationHandler"]) + table.grant_read_write_data(functions["ResponseHook"]) def _create_dynamodb_table(self) -> Table: table = dynamodb.Table( diff --git a/tests/e2e/idempotency/test_idempotency_dynamodb.py b/tests/e2e/idempotency/test_idempotency_dynamodb.py index ea4a319b76e..75d774e702e 100644 --- a/tests/e2e/idempotency/test_idempotency_dynamodb.py +++ b/tests/e2e/idempotency/test_idempotency_dynamodb.py @@ -41,6 +41,11 @@ def payload_tampering_validation_fn_arn(infrastructure: dict) -> str: return infrastructure.get("PayloadTamperingValidationHandlerArn", "") +@pytest.fixture +def response_hook_handler_fn_arn(infrastructure: dict) -> str: + return infrastructure.get("ResponseHookArn", "") + + @pytest.fixture def idempotency_table_name(infrastructure: dict) -> str: return infrastructure.get("DynamoDBTable", "") @@ -219,3 +224,29 @@ def test_payload_tampering_validation(payload_tampering_validation_fn_arn: str): lambda_arn=payload_tampering_validation_fn_arn, payload=json.dumps(tampered_transaction), ) + + +@pytest.mark.xdist_group(name="idempotency") +def test_response_hook_idempotency(response_hook_handler_fn_arn: str): + # GIVEN + payload = json.dumps({"message": "Powertools for AWS Lambda (Python)"}) + + # WHEN + # first execution + first_execution, _ = data_fetcher.get_lambda_response( + lambda_arn=response_hook_handler_fn_arn, + payload=payload, + ) + first_execution_response = first_execution["Payload"].read().decode("utf-8") + + # the second execution should include response hook + second_execution, _ = data_fetcher.get_lambda_response( + lambda_arn=response_hook_handler_fn_arn, + payload=payload, + ) + second_execution_response = second_execution["Payload"].read().decode("utf-8") + + # THEN first execution should not trigger response hook + # THEN seconde execution must trigger response hook + assert "x-response-hook" not in first_execution_response + assert "x-response-hook" in second_execution_response diff --git a/tests/e2e/idempotency_redis/handlers/response_hook.py b/tests/e2e/idempotency_redis/handlers/response_hook.py new file mode 100644 index 00000000000..b56d09f25b9 --- /dev/null +++ b/tests/e2e/idempotency_redis/handlers/response_hook.py @@ -0,0 +1,29 @@ +import os + +from aws_lambda_powertools.utilities.idempotency import ( + DynamoDBPersistenceLayer, + IdempotencyConfig, + idempotent, +) +from aws_lambda_powertools.utilities.idempotency.persistence.datarecord import ( + DataRecord, +) + +TABLE_NAME = os.getenv("IdempotencyTable", "") +persistence_layer = DynamoDBPersistenceLayer(table_name=TABLE_NAME) + + +def my_response_hook(response: dict, idempotent_data: DataRecord) -> dict: + # Return inserted Header data into the Idempotent Response + response["x-response-hook"] = idempotent_data.idempotency_key + + # Must return the response here + return response + + +config = IdempotencyConfig(response_hook=my_response_hook) + + +@idempotent(config=config, persistence_store=persistence_layer) +def lambda_handler(event, context): + return {"message": "first_response"} diff --git a/tests/e2e/idempotency_redis/test_idempotency_redis.py b/tests/e2e/idempotency_redis/test_idempotency_redis.py index 47b16760b82..ee5502b2dec 100644 --- a/tests/e2e/idempotency_redis/test_idempotency_redis.py +++ b/tests/e2e/idempotency_redis/test_idempotency_redis.py @@ -32,6 +32,11 @@ def optional_idempotency_key_fn_arn(infrastructure: dict) -> str: return infrastructure.get("OptionalIdempotencyKeyHandlerArn", "") +@pytest.fixture +def response_hook_handler_fn_arn(infrastructure: dict) -> str: + return infrastructure.get("ResponseHookArn", "") + + @pytest.mark.xdist_group(name="idempotency-redis") def test_ttl_caching_expiration_idempotency(ttl_cache_expiration_handler_fn_arn: str): # GIVEN @@ -181,3 +186,29 @@ def test_optional_idempotency_key(optional_idempotency_key_fn_arn: str): assert first_execution_response != second_execution_response assert first_execution_response != third_execution_response assert second_execution_response != third_execution_response + + +@pytest.mark.xdist_group(name="idempotency") +def test_response_hook_idempotency(response_hook_handler_fn_arn: str): + # GIVEN + payload = json.dumps({"message": "Powertools for AWS Lambda (Python)"}) + + # WHEN + # first execution + first_execution, _ = data_fetcher.get_lambda_response( + lambda_arn=response_hook_handler_fn_arn, + payload=payload, + ) + first_execution_response = first_execution["Payload"].read().decode("utf-8") + + # the second execution should include response hook + second_execution, _ = data_fetcher.get_lambda_response( + lambda_arn=response_hook_handler_fn_arn, + payload=payload, + ) + second_execution_response = second_execution["Payload"].read().decode("utf-8") + + # THEN first execution should not trigger response hook + # THEN seconde execution must trigger response hook + assert "x-response-hook" not in first_execution_response + assert "x-response-hook" in second_execution_response diff --git a/tests/functional/idempotency/_boto3/test_idempotency.py b/tests/functional/idempotency/_boto3/test_idempotency.py index 1d969dc19c1..35f82333e9c 100644 --- a/tests/functional/idempotency/_boto3/test_idempotency.py +++ b/tests/functional/idempotency/_boto3/test_idempotency.py @@ -1963,3 +1963,54 @@ def lambda_handler(event, context): stubber.assert_no_pending_responses() stubber.deactivate() + + +@pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}, {"use_local_cache": True}], indirect=True) +def test_idempotent_lambda_already_completed_response_hook_is_called_with_none( + idempotency_config: IdempotencyConfig, + persistence_store: DynamoDBPersistenceLayer, + lambda_apigw_event, + timestamp_future, + hashed_idempotency_key, + lambda_context, +): + """ + Test idempotent decorator where event with matching event key has already been successfully processed + """ + + def idempotent_response_hook(response: Any, idempotent_data: DataRecord) -> Any: + """Modify the response provided by adding a new key""" + new_response: dict = {} + new_response["idempotent_response"] = True + new_response["response"] = response + new_response["idempotent_expiration"] = idempotent_data.get_expiration_datetime() + + return new_response + + idempotency_config.response_hook = idempotent_response_hook + + stubber = stub.Stubber(persistence_store.client) + ddb_response = { + "Item": { + "id": {"S": hashed_idempotency_key}, + "expiration": {"N": timestamp_future}, + "data": {"S": "null"}, + "status": {"S": "COMPLETED"}, + }, + } + stubber.add_client_error("put_item", "ConditionalCheckFailedException", modeled_fields=ddb_response) + stubber.activate() + + @idempotent(config=idempotency_config, persistence_store=persistence_store) + def lambda_handler(event, context): + raise Exception + + lambda_resp = lambda_handler(lambda_apigw_event, lambda_context) + + # Then idempotent_response value will be added to the response + assert lambda_resp["idempotent_response"] + assert lambda_resp["response"] is None + assert lambda_resp["idempotent_expiration"] == datetime.datetime.fromtimestamp(int(timestamp_future)) + + stubber.assert_no_pending_responses() + stubber.deactivate()