fix(idempotency): capture the lambda context automatically

Author: rubenfonseca

aws_lambda_powertools/utilities/idempotency/base.py

@@ -130,13 +130,8 @@ def _get_remaining_time_in_millis(self) -> Optional[int]:
             Remaining time in millis, or None if the remaining time cannot be determined.
         """
 
-        # Look to see if we have stored a Lambda Context
         if self.config.lambda_context is not None:
-            self.config.lambda_context.get_remaining_time_in_millis()
-
-        # Look into fn_args to see if we have a lambda context
-        if self.fn_args and len(self.fn_args) == 2 and getattr(self.fn_args[1], "get_remaining_time_in_millis", None):
-            return self.fn_args[1].get_remaining_time_in_millis()
+            return self.config.lambda_context.get_remaining_time_in_millis()
 
         return None
 

aws_lambda_powertools/utilities/idempotency/idempotency.py

@@ -62,6 +62,8 @@ def idempotent(
         return handler(event, context)
 
     config = config or IdempotencyConfig()
+    config.register_lambda_context(context)
+
     args = event, context
     idempotency_handler = IdempotencyHandler(
         function=handler,

docs/utilities/idempotency.md

@@ -392,11 +392,11 @@ In cases where the [Lambda invocation
 expires](https://aws.amazon.com/premiumsupport/knowledge-center/lambda-verify-invocation-timeouts/),
 Powertools doesn't have the chance to set the idempotency record to `EXPIRED`.
 This means that the record would normally have been locked until `expire_seconds` have
-passsed.
+passed.
 
 However, when Powertools has access to the Lambda invocation context, we are able to calculate the remaining
 available time for the invocation, and save it on the idempotency record. This way, if a second invocation happens
-after this timestamp, and the record is still marked `INPROGRESS`, we execute the inovcation again as if it was
+after this timestamp, and the record is still marked `INPROGRESS`, we execute the invocation again as if it was
 already expired. This means that if an invocation expired during execution, it will be quickly executed again on the
 next retry.
 
@@ -435,9 +435,9 @@ sequenceDiagram
 </center>
 
 ???+ info "Info: Calculating the remaining available time"
-    When using the `idempotent` decorator, we captura and calculate the remaining available time for you.
+    When using the `idempotent` decorator, we capture and calculate the remaining available time for you.
     However, when using the `idempotent_function`, the functionality doesn't work out of the box. You'll
-    need to register the Lambda context on your handler:
+    need to register the Lambda context in your handler:
 
 ```python hl_lines="8 16" title="Registering the Lambda context"
 from aws_lambda_powertools.utilities.data_classes.sqs_event import SQSRecord

tests/functional/idempotency/conftest.py

@@ -1,6 +1,5 @@
 import datetime
 import json
-from collections import namedtuple
 from decimal import Decimal
 from unittest import mock
 
@@ -32,14 +31,17 @@ def lambda_apigw_event():
 
 @pytest.fixture
 def lambda_context():
-    lambda_context = {
-        "function_name": "test-func",
-        "memory_limit_in_mb": 128,
-        "invoked_function_arn": "arn:aws:lambda:eu-west-1:809313241234:function:test-func",
-        "aws_request_id": "52fdfc07-2182-154f-163f-5f0f9a621d72",
-    }
+    class LambdaContext:
+        def __init__(self):
+            self.function_name = "test-func"
+            self.memory_limit_in_mb = 128
+            self.invoked_function_arn = "arn:aws:lambda:eu-west-1:809313241234:function:test-func"
+            self.aws_request_id = "52fdfc07-2182-154f-163f-5f0f9a621d72"
+
+        def get_remaining_time_in_millis(self) -> int:
+            return 1000
 
-    return namedtuple("LambdaContext", lambda_context.keys())(*lambda_context.values())
+    return LambdaContext()
 
 
 @pytest.fixture
@@ -142,7 +144,12 @@ def expected_params_put_item(hashed_idempotency_key):
             "#in_progress_expiry": "in_progress_expiration",
         },
         "ExpressionAttributeValues": {":now": stub.ANY, ":inprogress": "INPROGRESS"},
-        "Item": {"expiration": stub.ANY, "id": hashed_idempotency_key, "status": "INPROGRESS"},
+        "Item": {
+            "expiration": stub.ANY,
+            "id": hashed_idempotency_key,
+            "status": "INPROGRESS",
+            "in_progress_expiration": stub.ANY,
+        },
         "TableName": "TEST_TABLE",
     }
 
@@ -163,6 +170,7 @@ def expected_params_put_item_with_validation(hashed_idempotency_key, hashed_vali
         "ExpressionAttributeValues": {":now": stub.ANY, ":inprogress": "INPROGRESS"},
         "Item": {
             "expiration": stub.ANY,
+            "in_progress_expiration": stub.ANY,
             "id": hashed_idempotency_key,
             "status": "INPROGRESS",
             "validation": hashed_validation_key,

tests/functional/idempotency/test_idempotency.py

@@ -768,37 +768,22 @@ def lambda_handler(event, context):
     stubber.deactivate()
 
 
-@pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}, {"use_local_cache": True}], indirect=True)
-def test_idempotent_lambda_expires_in_progress_unavailable_remaining_time(
-    idempotency_config: IdempotencyConfig,
-    persistence_store: DynamoDBPersistenceLayer,
-    lambda_apigw_event,
-    lambda_response,
-    lambda_context,
-    expected_params_put_item,
-    expected_params_update_item,
-):
-    stubber = stub.Stubber(persistence_store.table.meta.client)
-
-    ddb_response = {}
-    stubber.add_response("put_item", ddb_response, expected_params_put_item)
-    stubber.add_response("update_item", ddb_response, expected_params_update_item)
-
-    stubber.activate()
+def test_idempotent_lambda_expires_in_progress_unavailable_remaining_time():
+    mock_event = {"data": "value"}
+    idempotency_key = "test-func.function#" + hash_idempotency_key(mock_event)
+    persistence_layer = MockPersistenceLayer(expected_idempotency_key=idempotency_key)
+    expected_result = {"message": "Foo"}
 
-    @idempotent(config=idempotency_config, persistence_store=persistence_store)
-    def lambda_handler(event, context):
-        return lambda_response
+    @idempotent_function(persistence_store=persistence_layer, data_keyword_argument="record")
+    def function(record):
+        return expected_result
 
     with warnings.catch_warnings(record=True) as w:
         warnings.simplefilter("default")
-        lambda_handler(lambda_apigw_event, lambda_context)
+        function(record=mock_event)
         assert len(w) == 1
         assert str(w[-1].message) == "Expires in progress is enabled but we couldn't determine the remaining time left"
 
-    stubber.assert_no_pending_responses()
-    stubber.deactivate()
-
 
 def test_data_record_invalid_status_value():
     data_record = DataRecord("key", status="UNSUPPORTED_STATUS")

tests/functional/idempotency/utils.py

@@ -29,7 +29,12 @@ def build_idempotency_put_item_stub(
             "#in_progress_expiry": "in_progress_expiration",
         },
         "ExpressionAttributeValues": {":now": stub.ANY, ":inprogress": "INPROGRESS"},
-        "Item": {"expiration": stub.ANY, "id": idempotency_key_hash, "status": "INPROGRESS"},
+        "Item": {
+            "expiration": stub.ANY,
+            "id": idempotency_key_hash,
+            "status": "INPROGRESS",
+            "in_progress_expiration": stub.ANY,
+        },
         "TableName": "TEST_TABLE",
     }