Added some more tests + formatted code

rate_limiters/tests/test_fixed_window.py

@@ -2,6 +2,7 @@ import pytest
 
 from fixed_window import FixedWindowLimiter
 
+
 def test_limit_respected():
     limiter = FixedWindowLimiter(10, 10)
     try:
@@ -10,7 +11,8 @@ def test_limit_respected():
                 pass
     except RuntimeError as e:
         assert False, f"Limiter raised an exception -> {e}"
-    
+
+
 def test_limit_enforced():
     limiter = FixedWindowLimiter(10, 10)
     with pytest.raises(RuntimeError) as e:

rate_limiters/tests/test_leaky_bucket.py

@@ -5,11 +5,11 @@ import pytest
 
 from leaky_bucket import LeakyBucketLimiter
 
+
 @pytest.mark.asyncio
 async def test_duration() -> None:
     async def test_function():
         pass
-    
 
     limiter = LeakyBucketLimiter(10, 100)
     await limiter.start()
@@ -26,4 +26,4 @@ async def test_duration() -> None:
     assert duration < 2.1
 
     for item in results:
-        assert isinstance(item, Exception) is False
+        assert isinstance(item, Exception) is False

rate_limiters/tests/test_middleware_starlette.py

@@ -1,10 +1,15 @@
 from time import sleep
 
+from cachetools import TTLCache
+from starlette.applications import Starlette
 from starlette.testclient import TestClient
 
-from middleware_starlette import app
+from middleware_starlette import RateLimiterMiddleware, create_app, routes
+
 
 def test_app_runs():
+    app = create_app()
+
     client = TestClient(app)
 
     response = client.get("/")
@@ -14,6 +19,8 @@ def test_app_runs():
 
 
 def test_below_rate_limit():
+    app = create_app()
+
     client = TestClient(app)
 
     # Limit is 10/s
@@ -26,6 +33,8 @@ def test_below_rate_limit():
 
 
 def test_above_rate_limit():
+    app = create_app()
+
     client = TestClient(app)
 
     successful = 0
@@ -46,3 +55,96 @@ def test_above_rate_limit():
     assert successful != 0
     assert rate_limited != 0
 
+
+def find_middleware(stack, middleware_cls):
+    current = stack
+
+    while hasattr(current, "app"):
+        if isinstance(current, middleware_cls):
+            return current
+        current = current.app
+
+    return None
+
+
+def test_rate_limiter_cache():
+    app = create_app()
+
+    client = TestClient(app)
+
+    client.get("/")
+
+    middleware = find_middleware(
+        app.middleware_stack,
+        RateLimiterMiddleware,
+    )
+
+    assert middleware is not None
+
+    assert "testclient" in middleware.clients
+    assert len(middleware.clients) == 1
+
+
+def test_cache_ttl():
+    # Setup app with a short TTL
+    class ShortTTLLimiterMiddleware(RateLimiterMiddleware):
+        def __init__(self, app):
+            super().__init__(app)
+            self.clients = TTLCache(maxsize=10, ttl=2)  # Clear after 2 seconds
+
+    app = Starlette(routes=routes)
+    app.add_middleware(ShortTTLLimiterMiddleware)
+
+    client = TestClient(app)
+
+    client.get("/")
+
+    middleware = find_middleware(
+        app.middleware_stack,
+        RateLimiterMiddleware,
+    )
+
+    assert middleware is not None
+
+    client.get("/", headers={"X-Forwarded-For": "1.1.1.1"})
+
+    assert middleware.clients.get("1.1.1.1") is not None
+
+    sleep(2)
+
+    assert middleware.clients.get("1.1.1.1") is None
+
+
+def test_cache_eviction():
+    # Setup app with a tiny cache for testing
+    class TinyLimiterMiddleware(RateLimiterMiddleware):
+        def __init__(self, app):
+            super().__init__(app)
+            self.clients = TTLCache(maxsize=2, ttl=10)  # Only space for 2 IPs
+
+    app = Starlette(routes=routes)
+    app.add_middleware(TinyLimiterMiddleware)
+
+    client = TestClient(app)
+
+    client.get("/")
+
+    middleware = find_middleware(
+        app.middleware_stack,
+        RateLimiterMiddleware,
+    )
+
+    assert middleware is not None
+
+    client.get("/", headers={"X-Forwarded-For": "1.1.1.1"})
+    client.get("/", headers={"X-Forwarded-For": "2.2.2.2"})
+
+    assert middleware.clients.get("1.1.1.1") is not None
+    assert middleware.clients.get("2.2.2.2") is not None
+    assert middleware.clients.get("3.3.3.3") is None
+
+    client.get("/", headers={"X-Forwarded-For": "3.3.3.3"})
+
+    assert middleware.clients.get("1.1.1.1") is None
+    assert middleware.clients.get("2.2.2.2") is not None
+    assert middleware.clients.get("3.3.3.3") is not None

rate_limiters/tests/test_token_bucket.py

@@ -1,6 +1,9 @@
-import pytest
 from time import sleep
-from token_bucket import TokenBucketLimiter, RateLimitExceeded
+
+import pytest
+
+from token_bucket import RateLimitExceeded, TokenBucketLimiter
+
 
 def test_limit_respected() -> None:
     limiter = TokenBucketLimiter(10, 10)
@@ -20,6 +23,5 @@ def test_limit_exceeded() -> None:
         for i in range(20):
             with limiter:
                 continue
-    
-    assert "Rate exceeded" in str(e.value)
 
+    assert "Rate exceeded" in str(e.value)