Livrare LOT 1 - Didi

ai_platform/modules/llm-inference/tests/test_concurrency.py

@@ -0,0 +1,139 @@
+"""Tests for concurrency limiter functionality."""
+
+import asyncio
+
+import pytest
+from fastapi import HTTPException
+
+from llm_inference.api.dependencies import ConcurrencyLimiter
+
+
+class TestConcurrencyLimiter:
+    """Tests for ConcurrencyLimiter class."""
+
+    @pytest.mark.asyncio
+    async def test_acquire_within_limit(self) -> None:
+        """Test acquiring slots within concurrency limit."""
+        limiter = ConcurrencyLimiter(max_concurrent=2)
+
+        async with limiter.acquire():
+            assert limiter.current_count == 1
+            assert limiter.available == 1
+
+        assert limiter.current_count == 0
+        assert limiter.available == 2
+
+    @pytest.mark.asyncio
+    async def test_acquire_exceeds_limit_raises_503(self) -> None:
+        """Test that exceeding limit raises 503 HTTPException."""
+        limiter = ConcurrencyLimiter(max_concurrent=1)
+
+        async with limiter.acquire():
+            # Try to acquire another slot while one is held
+            with pytest.raises(HTTPException) as exc_info:
+                async with limiter.acquire():
+                    pass
+
+            assert exc_info.value.status_code == 503
+            assert "Too many concurrent requests" in str(exc_info.value.detail)
+
+    @pytest.mark.asyncio
+    async def test_acquire_blocking_mode(self) -> None:
+        """Test blocking mode waits for available slot."""
+        limiter = ConcurrencyLimiter(max_concurrent=1)
+        results: list[int] = []
+
+        async def task(task_id: int) -> None:
+            async with limiter.acquire(blocking=True):
+                results.append(task_id)
+                await asyncio.sleep(0.01)
+
+        # Start multiple tasks - they should execute sequentially
+        await asyncio.gather(task(1), task(2), task(3))
+
+        # All tasks should complete (order may vary due to concurrency)
+        assert sorted(results) == [1, 2, 3]
+
+    @pytest.mark.asyncio
+    async def test_concurrent_acquire_releases_properly(self) -> None:
+        """Test that slots are properly released even with concurrent access."""
+        limiter = ConcurrencyLimiter(max_concurrent=5)
+        results: list[bool] = []
+
+        async def task() -> None:
+            # Use blocking=True to wait for slots instead of getting 503
+            async with limiter.acquire(blocking=True):
+                results.append(True)
+                await asyncio.sleep(0.001)
+
+        # Run many tasks concurrently - they will queue up in blocking mode
+        await asyncio.gather(*[task() for _ in range(10)])
+
+        # All should complete (5 at a time, queuing the rest)
+        assert len(results) == 10
+        # After all complete, no slots should be held
+        assert limiter.current_count == 0
+
+    @pytest.mark.asyncio
+    async def test_acquire_releases_on_exception(self) -> None:
+        """Test that slot is released even if body raises exception."""
+        limiter = ConcurrencyLimiter(max_concurrent=1)
+
+        with pytest.raises(ValueError):
+            async with limiter.acquire():
+                assert limiter.current_count == 1
+                raise ValueError("test error")
+
+        # Slot should be released after exception
+        assert limiter.current_count == 0
+        assert limiter.available == 1
+
+    @pytest.mark.asyncio
+    async def test_counter_lock_prevents_race_condition(self) -> None:
+        """Test that counter updates are atomic with lock."""
+        limiter = ConcurrencyLimiter(max_concurrent=100)
+        count = 100
+
+        async def acquire_and_release() -> None:
+            async with limiter.acquire():
+                await asyncio.sleep(0.001)
+
+        # Run many concurrent acquires/releases
+        await asyncio.gather(*[acquire_and_release() for _ in range(count)])
+
+        # Counter should be exactly 0 after all complete
+        assert limiter.current_count == 0
+
+    @pytest.mark.asyncio
+    async def test_properties_accurate_during_use(self) -> None:
+        """Test that current_count and available properties are accurate."""
+        limiter = ConcurrencyLimiter(max_concurrent=3)
+
+        assert limiter.current_count == 0
+        assert limiter.available == 3
+
+        async with limiter.acquire():
+            assert limiter.current_count == 1
+            assert limiter.available == 2
+
+            async with limiter.acquire():
+                assert limiter.current_count == 2
+                assert limiter.available == 1
+
+            assert limiter.current_count == 1
+
+        assert limiter.current_count == 0
+
+    @pytest.mark.asyncio
+    async def test_retry_after_header_in_503_response(self) -> None:
+        """Test that 503 response includes Retry-After header."""
+        limiter = ConcurrencyLimiter(max_concurrent=1)
+
+        async with limiter.acquire():
+            with pytest.raises(HTTPException) as exc_info:
+                async with limiter.acquire():
+                    pass
+
+            assert exc_info.value.headers is not None
+            assert "Retry-After" in exc_info.value.headers
+            assert exc_info.value.headers["Retry-After"] == "5"