fix(idempotency): fix response hook invocation when function returns None (#5251)

Fix reponse hook invocation with None

Author: leandrodamascena

Diff for: 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 for: 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 for: 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 for: 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 for: 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 for: 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 for: 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()