use in-memory cache, add tests

aggrec/key_cache.py

@@ -3,10 +3,15 @@
 from abc import abstractmethod
 
 import redis
+from expiringdict import ExpiringDict
 
 logger = logging.getLogger(__name__)
 
 
+DEFAULT_MEMORY_CACHE_SIZE = 1000
+DEFAULT_MEMORY_CACHE_TTL = 60
+
+
 class KeyCache:
     @abstractmethod
     def get(self, key: str) -> bytes | None:
@@ -17,16 +22,30 @@ def set(self, key: str, value: bytes, ttl: int | None = None) -> None:
         pass
 
 
-class NoyKeyCache(KeyCache):
+class DummyKeyCache(KeyCache):
     def get(self, key: str) -> bytes | None:
         return None
 
     def set(self, key: str, value: bytes, ttl: int | None = None) -> None:
         pass
 
 
+class MemoryKeyCache(KeyCache):
+    def __init__(self, size: int = DEFAULT_MEMORY_CACHE_SIZE, ttl: int = DEFAULT_MEMORY_CACHE_TTL):
+        self.cache = ExpiringDict(max_len=size, max_age_seconds=ttl)
+
+    def get(self, key: str) -> bytes | None:
+        res = self.cache.get(key)
+        logger.debug("Cache GET %s (%s)", key, "hit" if res else "miss")
+        return res
+
+    def set(self, key: str, value: bytes, ttl: int | None = None) -> None:
+        logger.debug("Cache SET %s", key)
+        self.cache[key] = value
+
+
 class RedisKeyCache(KeyCache):
-    def __init__(self, redis_client: redis.Redis, default_ttl: int | None = None):
+    def __init__(self, redis_client: redis.Redis, default_ttl: int):
         self.redis_client = redis_client
         self.default_ttl = default_ttl
 

aggrec/server.py

@@ -14,7 +14,7 @@
 import aggrec.extras
 
 from . import OPENAPI_METADATA, __verbose_version__
-from .key_cache import NoyKeyCache, RedisKeyCache
+from .key_cache import MemoryKeyCache, RedisKeyCache
 from .logging import JsonFormatter  # noqa
 from .settings import Settings
 from .telemetry import configure_opentelemetry
@@ -75,7 +75,7 @@ def __init__(self, settings: Settings):
             self.logger.debug("Using REDIS at %s:%d", self.settings.redis.host, self.settings.redis.port)
             self.key_cache = RedisKeyCache(redis_client=redis_client, default_ttl=self.settings.redis.ttl)
         else:
-            self.key_cache = NoyKeyCache()
+            self.key_cache = MemoryKeyCache()
 
     @staticmethod
     def connect_mongodb(settings: Settings):

pyproject.toml

@@ -31,13 +31,15 @@ opentelemetry-instrumentation-pymongo = "^0.48b0"
 http-sf = "^1.0.2"
 httpx = "^0.27.2"
 redis = "^5.1.1"
+expiringdict = "^1.2.2"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.2.0"
 ruff = "^0.6.3"
 pytest-ruff = "^0.4.1"
 pytest-asyncio = "^0.24.0"
 fakeredis = "^2.25.1"
+pytest-httpx = "^0.32.0"
 
 [build-system]
 requires = ["poetry-core"]

tests/test_key_cache.py

@@ -0,0 +1,42 @@
+import fakeredis
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import ed25519
+
+from aggrec.key_cache import MemoryKeyCache, RedisKeyCache
+
+
+def test_redis_cache():
+    key_id = "xyzzy"
+    public_key = ed25519.Ed25519PrivateKey.generate().public_key()
+    public_key_pem = public_key.public_bytes(
+        encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo
+    )
+
+    redis_client = fakeredis.FakeRedis()
+    key_cache = RedisKeyCache(redis_client=redis_client, default_ttl=60)
+
+    res = key_cache.get(key_id)
+    assert res is None
+
+    key_cache.set(key_id, public_key_pem)
+
+    res = key_cache.get(key_id)
+    assert res == public_key_pem
+
+
+def test_memory_cache():
+    key_id = "xyzzy"
+    public_key = ed25519.Ed25519PrivateKey.generate().public_key()
+    public_key_pem = public_key.public_bytes(
+        encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo
+    )
+
+    key_cache = MemoryKeyCache()
+
+    res = key_cache.get(key_id)
+    assert res is None
+
+    key_cache.set(key_id, public_key_pem)
+
+    res = key_cache.get(key_id)
+    assert res == public_key_pem

tests/test_key_resolver.py

@@ -0,0 +1,21 @@
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import ed25519
+from pytest_httpx import HTTPXMock
+
+from aggrec.key_cache import MemoryKeyCache
+from aggrec.key_resolver import UrlKeyResolver
+
+
+def test_url_key_resolver(httpx_mock: HTTPXMock):
+    key_id = "xyzzy"
+    public_key = ed25519.Ed25519PrivateKey.generate().public_key()
+    public_key_pem = public_key.public_bytes(
+        encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo
+    )
+
+    httpx_mock.add_response(url=f"https://keys/{key_id}.pem", content=public_key_pem)
+
+    resolver = UrlKeyResolver(client_database_base_url="https://keys", key_cache=MemoryKeyCache())
+
+    res = resolver.resolve_public_key(key_id)
+    assert res == public_key