Livrare LOT 1 - Didi

ai_platform/modules/web/tests/test_middleware.py

@@ -0,0 +1,44 @@
+"""Tests for rate limiting middleware."""
+
+import pytest
+
+from web.api.middleware import TokenBucket
+
+
+class TestTokenBucket:
+    """Tests for TokenBucket rate limiter."""
+
+    @pytest.mark.asyncio
+    async def test_initial_burst(self) -> None:
+        """Test that initial burst tokens are available."""
+        bucket = TokenBucket(rate=10.0, burst=5)
+
+        # Should be able to acquire burst tokens
+        for _ in range(5):
+            allowed, _ = await bucket.acquire()
+            assert allowed is True
+
+        # Next one should fail
+        allowed, _ = await bucket.acquire()
+        assert allowed is False
+
+    @pytest.mark.asyncio
+    async def test_acquire_returns_retry_after_when_empty(self) -> None:
+        """Test that acquire returns positive retry_after when empty."""
+        bucket = TokenBucket(rate=1.0, burst=1)
+
+        # Drain the bucket
+        await bucket.acquire()
+
+        allowed, retry_after = await bucket.acquire()
+        assert allowed is False
+        assert retry_after > 0.0
+
+    @pytest.mark.asyncio
+    async def test_acquire_returns_zero_retry_when_available(self) -> None:
+        """Test that acquire returns 0 retry_after when tokens available."""
+        bucket = TokenBucket(rate=10.0, burst=5)
+
+        allowed, retry_after = await bucket.acquire()
+        assert allowed is True
+        assert retry_after == 0.0