44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
|
|
from app.llm.queue import LLMQueue, QueueFullError
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_queue_runs_calls_up_to_concurrency_limit():
|
|
queue = LLMQueue(max_concurrency=2, max_queue_depth=10)
|
|
concurrent_count = 0
|
|
max_observed = 0
|
|
|
|
async def slow_call():
|
|
nonlocal concurrent_count, max_observed
|
|
concurrent_count += 1
|
|
max_observed = max(max_observed, concurrent_count)
|
|
await asyncio.sleep(0.05)
|
|
concurrent_count -= 1
|
|
return "done"
|
|
|
|
results = await asyncio.gather(*(queue.submit(slow_call) for _ in range(5)))
|
|
|
|
assert results == ["done"] * 5
|
|
assert max_observed == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_queue_raises_when_depth_exceeded():
|
|
queue = LLMQueue(max_concurrency=1, max_queue_depth=1)
|
|
|
|
async def slow_call():
|
|
await asyncio.sleep(0.1)
|
|
return "done"
|
|
|
|
task1 = asyncio.create_task(queue.submit(slow_call))
|
|
task2 = asyncio.create_task(queue.submit(slow_call))
|
|
await asyncio.sleep(0.01) # let task1 start running, task2 start waiting
|
|
|
|
with pytest.raises(QueueFullError):
|
|
await queue.submit(slow_call)
|
|
|
|
await asyncio.gather(task1, task2)
|