testing rate-limiting requests via ASGI middleware in starlette

rate_limiters/tests/test_leaky_bucket.py

@@ -0,0 +1,29 @@
+import asyncio
+from time import time
+
+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()
+
+    start = time()
+
+    tasks = [asyncio.create_task(limiter.execute(test_function)) for i in range(20)]
+
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+    end = time()
+    duration = end - start
+
+    assert duration < 2.1
+
+    for item in results:
+        assert isinstance(item, Exception) is False

rate_limiters/tests/test_middleware_starlette.py

@@ -0,0 +1,49 @@
+from time import sleep
+
+from starlette.testclient import TestClient
+
+from middleware_starlette import app
+
+def test_app_runs():
+    client = TestClient(app)
+
+    response = client.get("/")
+
+    assert response.status_code == 200
+    assert "Hello there" in response.text
+
+
+def test_below_rate_limit():
+    client = TestClient(app)
+
+    # Limit is 10/s
+    for i in range(20):
+        response = client.get("/")
+
+        assert response.status_code == 200
+        assert "Hello there" in response.text
+        sleep(0.1)
+
+
+def test_above_rate_limit():
+    client = TestClient(app)
+
+    successful = 0
+    rate_limited = 0
+
+    # Limit is 10/s
+    for i in range(30):
+        response = client.get("/")
+
+        if response.status_code != 200:
+            rate_limited += 1
+            assert response.status_code == 429
+            assert "Rate limit exceeded" in response.text
+            assert response.headers.get("retry-after") is not None
+        else:
+            successful += 1
+
+    assert successful != 0
+    assert rate_limited != 0
+
+