fix(idempotent): Correctly raise IdempotencyKeyError (#378)

Author: Michael Brewer

aws_lambda_powertools/utilities/idempotency/idempotency.py

@@ -11,6 +11,7 @@
     IdempotencyInconsistentStateError,
     IdempotencyItemAlreadyExistsError,
     IdempotencyItemNotFoundError,
+    IdempotencyKeyError,
     IdempotencyPersistenceLayerError,
     IdempotencyValidationError,
 )
@@ -132,6 +133,8 @@ def handle(self) -> Any:
             # We call save_inprogress first as an optimization for the most common case where no idempotent record
             # already exists. If it succeeds, there's no need to call get_record.
             self.persistence_store.save_inprogress(event=self.event, context=self.context)
+        except IdempotencyKeyError:
+            raise
         except IdempotencyItemAlreadyExistsError:
             # Now we know the item already exists, we can retrieve it
             record = self._get_idempotency_record()

tests/functional/idempotency/test_idempotency.py

@@ -828,3 +828,20 @@ def lambda_handler(event, context):
     stubber.assert_no_pending_responses()
     stubber.deactivate()
     assert "Failed to save in progress record to idempotency store" == e.value.args[0]
+
+
+def test_handler_raise_idempotency_key_error(persistence_store: DynamoDBPersistenceLayer, lambda_context):
+    # GIVEN raise_on_no_idempotency_key is True
+    idempotency_config = IdempotencyConfig(event_key_jmespath="idemKey", raise_on_no_idempotency_key=True)
+
+    # WHEN handling the idempotent call
+    # AND save_inprogress raises a IdempotencyKeyError
+    @idempotent(persistence_store=persistence_store, config=idempotency_config)
+    def handler(event, context):
+        raise ValueError("Should not be raised")
+
+    # THEN idempotent should re-raise the IdempotencyKeyError
+    with pytest.raises(IdempotencyKeyError) as e:
+        handler({}, lambda_context)
+
+    assert "No data found to create a hashed idempotency_key" == e.value.args[0]