feat(idempotency): Fix KeyError when local_cache is True and an error is raised in the lambda handler (#300)

Author: Michael Brewer

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -272,7 +272,8 @@ def _retrieve_from_cache(self, idempotency_key: str):
     def _delete_from_cache(self, idempotency_key: str):
         if not self.use_local_cache:
             return
-        del self._cache[idempotency_key]
+        if idempotency_key in self._cache:
+            del self._cache[idempotency_key]
 
     def save_success(self, event: Dict[str, Any], result: dict) -> None:
         """

tests/functional/idempotency/test_idempotency.py

@@ -627,3 +627,14 @@ def test_user_local_disabled(persistence_store):
     # THEN raise AttributeError
     # AND don't have a _cache attribute
     assert not hasattr("persistence_store", "_cache")
+
+
+@pytest.mark.parametrize("persistence_store", [{"use_local_cache": True}], indirect=True)
+def test_delete_from_cache_when_empty(persistence_store):
+    # GIVEN use_local_cache is True AND the local cache is empty
+    try:
+        # WHEN we _delete_from_cache
+        persistence_store._delete_from_cache("key_does_not_exist")
+    except KeyError:
+        # THEN we should not get a KeyError
+        pytest.fail("KeyError should not happen")

tests/unit/test_lru_cache.py

@@ -56,3 +56,25 @@ def test_setitem_moves_to_end(populated_cache):
 
     assert last_item == f"key_{random_value}"
     assert populated_cache[f"key_{random_value}"] == f"new_val_{random_value}"
+
+
+def test_lru_pop_failing():
+    cache = LRUDict()
+    key = "test"
+    cache[key] = "value"
+    try:
+        cache.pop(key, None)
+        pytest.fail("GitHub #300: LRUDict pop bug has been fixed :)")
+    except KeyError as e:
+        assert e.args[0] == key
+
+
+def test_lru_del():
+    cache = LRUDict()
+    key = "test"
+    cache[key] = "value"
+    assert len(cache) == 1
+    if key in cache:
+        del cache[key]
+    assert key not in cache
+    assert len(cache) == 0