139 lines
5 KiB
Python
139 lines
5 KiB
Python
"""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"
|