add mock redis, redis validation, sentinel support

roger-zhangg · roger-zhangg · commit ad26e31ce84b · 2023-08-24T00:03:37.000Z
diff --git a/Makefile b/Makefile
@@ -40,6 +40,8 @@ test-idempotency-redis:
 	docker run --name test-idempotency-redis -d -p 63005:6379 redis
 	poetry run pytest tests/integration/idempotency;docker stop test-idempotency-redis;docker rm test-idempotency-redis
 
+
+
 e2e-test:
 	python parallel_run_e2e.py
 
diff --git a/aws_lambda_powertools/utilities/idempotency/exceptions.py b/aws_lambda_powertools/utilities/idempotency/exceptions.py
@@ -71,3 +71,9 @@ class IdempotencyKeyError(BaseError):
     """
     Payload does not contain an idempotent key
     """
+
+
+class IdempotencyRedisClientConfigError(BaseError):
+    """
+    Redis Client passed into persistant layer is not valid
+    """
diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/redis.py b/aws_lambda_powertools/utilities/idempotency/persistence/redis.py
@@ -1,16 +1,20 @@
 import datetime
 import logging
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 try:
     import redis  # type:ignore
 except ImportError:
     redis = None
 
+
+import redis
+
 from aws_lambda_powertools.utilities.idempotency import BasePersistenceLayer
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyItemAlreadyExistsError,
     IdempotencyItemNotFoundError,
+    IdempotencyRedisClientConfigError,
 )
 from aws_lambda_powertools.utilities.idempotency.persistence.base import (
     STATUS_CONSTANTS,
@@ -23,7 +27,7 @@
 class RedisCachePersistenceLayer(BasePersistenceLayer):
     def __init__(
         self,
-        connection,
+        connection: Union[redis.Redis, redis.cluster.RedisCluster],
         in_progress_expiry_attr: str = "in_progress_expiration",
         status_attr: str = "status",
         data_attr: str = "data",
@@ -44,6 +48,12 @@ def __init__(
         """
 
         # Initialize connection with Redis
+
+        if not hasattr(connection, "get_connection_kwargs"):
+            raise IdempotencyRedisClientConfigError
+        if not connection.get_connection_kwargs().get("decode_responses", False):
+            # Requires decode_responses to be true
+            raise IdempotencyRedisClientConfigError
         self._connection = connection
 
         self.in_progress_expiry_attr = in_progress_expiry_attr
@@ -106,29 +116,34 @@ def _put_record(self, data_record: DataRecord) -> None:
             #    - first time that this invocation key is used
             #    - previous invocation with the same key was deleted due to TTL
             idempotency_record = self._connection.hgetall(data_record.idempotency_key)
+            print(idempotency_record)
             if len(idempotency_record) > 0:
                 # record already exists.
 
                 # status is completed, so raise exception because it exists and still valid
                 if idempotency_record[self.status_attr] == STATUS_CONSTANTS["COMPLETED"]:
-                    raise
+                    raise IdempotencyItemAlreadyExistsError
 
                 # checking if in_progress_expiry_attr exists
                 # if in_progress_expiry_attr exist, must be lower than now
                 if self.in_progress_expiry_attr in idempotency_record and int(
                     idempotency_record[self.in_progress_expiry_attr],
                 ) > int(now.timestamp() * 1000):
-                    raise
+                    raise IdempotencyItemAlreadyExistsError
 
             logger.debug(f"Putting record on Redis for idempotency key: {data_record.idempotency_key}")
             self._connection.hset(**item)
             # hset type must set expiration after adding the record
             # Need to review this to get ttl in seconds
             # Q: should we replace self.expires_after_seconds with _get_expiry_timestamp? more consistent
             self._connection.expire(name=data_record.idempotency_key, time=self.expires_after_seconds)
-        except Exception:
-            logger.debug(f"Failed to put record for already existing idempotency key: {data_record.idempotency_key}")
-            raise IdempotencyItemAlreadyExistsError
+        except redis.exceptions.RedisError:
+            raise redis.exceptions.RedisError
+        except redis.exceptions.RedisClusterException:
+            raise redis.exceptions.RedisClusterException
+        except Exception as e:
+            logger.debug(f"encountered non-redis exception:{e}")
+            raise e
 
     def _update_record(self, data_record: DataRecord) -> None:
         item: Dict[str, Any] = {}
diff --git a/tests/functional/idempotency/persistence/test_redis_layer.py b/tests/functional/idempotency/persistence/test_redis_layer.py
@@ -0,0 +1,226 @@
+import copy
+import time as t
+
+import pytest
+
+from aws_lambda_powertools.utilities.idempotency import RedisCachePersistenceLayer
+from aws_lambda_powertools.utilities.idempotency.exceptions import (
+    IdempotencyAlreadyInProgressError,
+    IdempotencyItemAlreadyExistsError,
+    IdempotencyItemNotFoundError,
+    IdempotencyRedisClientConfigError,
+)
+from aws_lambda_powertools.utilities.idempotency.idempotency import (
+    idempotent,
+    idempotent_function,
+)
+
+
+@pytest.fixture
+def lambda_context():
+    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 LambdaContext()
+
+
+class MockRedis:
+    def __init__(self, decode_responses, cache, **kwargs):
+        self.cache = cache or {}
+        self.expire_dict = {}
+        self.decode_responses = decode_responses
+        self.acl = {}
+        self.username = ""
+
+    def hset(self, name, mapping):
+        self.expire_dict.pop(name, {})
+        self.cache[name] = mapping
+
+    # not covered by test yet.
+    def expire(self, name, time):
+        self.expire_dict[name] = t.time() + time
+
+    # return {} if no match
+    def hgetall(self, name):
+        if self.expire_dict.get(name, t.time() + 1) < t.time():
+            self.cache.pop(name, {})
+        return self.cache.get(name, {})
+
+    def get_connection_kwargs(self):
+        return {"decode_responses": self.decode_responses}
+
+    def auth(self, username, **kwargs):
+        self.username = username
+
+    def delete(self, name):
+        self.cache.pop(name, {})
+
+
+@pytest.fixture
+def persistence_store_standalone_redis():
+    # you will need to handle yourself the connection to pass again the password
+    # and avoid AuthenticationError at redis queries
+    redis_client = MockRedis(
+        host="localhost",
+        port="63005",
+        decode_responses=True,
+    )
+    return RedisCachePersistenceLayer(connection=redis_client)
+
+
+# test basic
+def test_idempotent_function_and_lambda_handler_redis_basic(
+    # idempotency_config: IdempotencyConfig,
+    persistence_store_standalone_redis: RedisCachePersistenceLayer,
+    lambda_context,
+):
+    mock_event = {"data": "value"}
+    persistence_layer = persistence_store_standalone_redis
+    expected_result = {"message": "Foo"}
+
+    @idempotent_function(persistence_store=persistence_layer, data_keyword_argument="record")
+    def record_handler(record):
+        return expected_result
+
+    @idempotent(persistence_store=persistence_layer)
+    def lambda_handler(event, context):
+        return expected_result
+
+    # WHEN calling the function
+    fn_result = record_handler(record=mock_event)
+    # WHEN calling lambda handler
+    handler_result = lambda_handler(mock_event, lambda_context)
+    # THEN we expect the function and lambda handler to execute successfully
+    assert fn_result == expected_result
+    assert handler_result == expected_result
+
+
+def test_idempotent_lambda_redis_no_decode():
+    redis_client = MockRedis(
+        host="localhost",
+        port="63005",
+        decode_responses=False,
+    )
+    # decode_responses=False will not be accepted
+    with pytest.raises(IdempotencyRedisClientConfigError):
+        RedisCachePersistenceLayer(connection=redis_client)
+
+
+def test_idempotent_function_and_lambda_handler_redis_cache(
+    persistence_store_standalone_redis: RedisCachePersistenceLayer,
+    lambda_context,
+):
+    mock_event = {"data": "value2"}
+    persistence_layer = persistence_store_standalone_redis
+    result = {"message": "Foo"}
+    expected_result = copy.deepcopy(result)
+
+    @idempotent_function(persistence_store=persistence_layer, data_keyword_argument="record")
+    def record_handler(record):
+        return result
+
+    @idempotent(persistence_store=persistence_layer)
+    def lambda_handler(event, context):
+        return result
+
+    # WHEN calling the function
+    fn_result = record_handler(record=mock_event)
+    # WHEN calling lambda handler
+    handler_result = lambda_handler(mock_event, lambda_context)
+    # THEN we expect the function and lambda handler to execute successfully
+    assert fn_result == expected_result
+    assert handler_result == expected_result
+
+    # modify the return to check if idem cache works
+    result = {"message": "Bar"}
+    fn_result2 = record_handler(record=mock_event)
+    # Second time calling lambda handler, test if same result
+    handler_result2 = lambda_handler(mock_event, lambda_context)
+    assert fn_result2 == expected_result
+    assert handler_result2 == expected_result
+
+    # modify the mock event to check if we got updated result
+    mock_event = {"data": "value3"}
+    fn_result3 = record_handler(record=mock_event)
+    # thrid time calling lambda handler, test if result updated
+    handler_result3 = lambda_handler(mock_event, lambda_context)
+    assert fn_result3 == result
+    assert handler_result3 == result
+
+
+# test idem-inprogress
+def test_idempotent_lambda_redis_in_progress(
+    persistence_store_standalone_redis: RedisCachePersistenceLayer,
+    lambda_context,
+):
+    """
+    Test idempotent decorator where lambda_handler is already processing an event with matching event key
+    """
+
+    mock_event = {"data": "value4"}
+    persistence_store = persistence_store_standalone_redis
+    lambda_response = {"foo": "bar"}
+
+    @idempotent(persistence_store=persistence_store)
+    def lambda_handler(event, context):
+        return lambda_response
+
+    # register the context first
+    lambda_handler(mock_event, lambda_context)
+    # save additional to in_progress
+    mock_event = {"data": "value7"}
+    try:
+        persistence_store.save_inprogress(mock_event, 1000)
+    except IdempotencyItemAlreadyExistsError:
+        pass
+
+    with pytest.raises(IdempotencyAlreadyInProgressError):
+        lambda_handler(mock_event, lambda_context)
+
+
+# test -remove
+def test_idempotent_lambda_redis_delete(
+    persistence_store_standalone_redis: RedisCachePersistenceLayer,
+    lambda_context,
+):
+    mock_event = {"data": "test_delete"}
+    persistence_layer = persistence_store_standalone_redis
+    result = {"message": "Foo"}
+
+    @idempotent(persistence_store=persistence_layer)
+    def lambda_handler(event, _):
+        return result
+
+    handler_result = lambda_handler(mock_event, lambda_context)
+    assert handler_result == result
+
+    # delete the idem and handler should output new result
+    persistence_layer.delete_record(mock_event, IdempotencyItemNotFoundError)
+    result = {"message": "Foo2"}
+    handler_result2 = lambda_handler(mock_event, lambda_context)
+    assert handler_result2 == result
+
+
+"""def test_idempotent_lambda_redis_credential(lambda_context):
+    redis_client = MockRedis(
+                host='localhost',
+                port='63005',
+                decode_responses=True,
+                )
+    pwd = "terriblePassword"
+    usr = "test_acl_denial"
+    redis_client.acl_setuser(username=usr, enabled=True, passwords="+"+pwd,keys='*',commands=['+hgetall','-set'])
+    redis_client.auth(password=pwd,username=usr)
+
+    @idempotent(persistence_store=RedisCachePersistenceLayer(connection=redis_client))
+    def lambda_handler(event, _):
+        return True
+    with pytest.raises(IdempotencyPersistenceLayerError):
+        handler_result = lambda_handler("test_Acl", lambda_context)"""
diff --git a/tests/integration/idempotency/setup_sentinel.sh b/tests/integration/idempotency/setup_sentinel.sh
@@ -0,0 +1,7 @@
+docker run --name redis_master -p 6379:6379 -d redis
+docker run --name redis_slave_1 -p 6380:6380  --link redis_master:redis_master -d redis redis-server  --slaveof redis_master 6379
+docker run --name redis_slave_2 -p 6381:6381 --link redis_master:redis_master -d redis redis-server  --slaveof redis_master 6379
+docker run --name redis_slave_3 -p 6382:6382  --link redis_master:redis_master -d redis redis-server  --slaveof redis_master 6379
+docker run --name redis_sentinel_1 -d -e REDIS_MASTER_HOST=redis_master -e REDIS_SENTINEL_PORT_NUMBER=26379 -e REDIS_SENTINEL_QUORUM=2 -p 26379:26379  --link redis_master:redis_master bitnami/redis-sentinel:latest
+docker run --name redis_sentinel_2 -d -e REDIS_MASTER_HOST=redis_master -e REDIS_SENTINEL_PORT_NUMBER=26380 -e REDIS_SENTINEL_QUORUM=2 -p 26380:26380  --link redis_master:redis_master bitnami/redis-sentinel:latest
+docker run --name redis_sentinel_3 -d -e REDIS_MASTER_HOST=redis_master -e REDIS_SENTINEL_PORT_NUMBER=26381 -e REDIS_SENTINEL_QUORUM=2 -p 26381:26381 --link redis_master:redis_master bitnami/redis-sentinel:latest
diff --git a/tests/integration/idempotency/test_idempotency_redis.py b/tests/integration/idempotency/test_idempotency_redis.py
