pablan/backend/tests/test_llm_gate.py

"""Several people asking at once, against an endpoint with a few slots.

The gate is the only place that bounds how much work reaches an endpoint, so
these tests are about the three answers it can give: go, wait, or busy.
"""

import asyncio
from collections.abc import Iterator

import pytest

from app.config import get_settings
from app.llm import gate
from app.llm.errors import LLMError

ENDPOINT = "http://endpoint.test/v1"


@pytest.fixture(autouse=True)
def _fresh_gate() -> Iterator[None]:
    gate.reset()
    yield
    gate.reset()


def _limits(
    monkeypatch: pytest.MonkeyPatch,
    *,
    parallel: int,
    wait: float = 20.0,
    queued: int = 24,
) -> None:
    settings = get_settings()
    monkeypatch.setattr(settings, "llm_max_parallel", parallel)
    monkeypatch.setattr(settings, "llm_queue_wait_seconds", wait)
    monkeypatch.setattr(settings, "llm_max_queued", queued)
    gate.reset()


async def test_only_as_many_calls_reach_the_endpoint_as_it_has_slots(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _limits(monkeypatch, parallel=2)
    concurrent = 0
    peak = 0
    release = asyncio.Event()

    async def call() -> None:
        nonlocal concurrent, peak
        async with gate.slot(ENDPOINT, "chat"):
            concurrent += 1
            peak = max(peak, concurrent)
            await release.wait()
            concurrent -= 1

    tasks = [asyncio.create_task(call()) for _ in range(6)]
    await asyncio.sleep(0)  # let everyone reach the gate
    assert peak == 2, "more calls were in flight than the endpoint has slots"
    release.set()
    await asyncio.gather(*tasks)
    assert peak == 2


async def test_a_waiting_call_gets_the_slot_the_previous_one_frees(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _limits(monkeypatch, parallel=1)
    order: list[str] = []
    first_in = asyncio.Event()
    let_go = asyncio.Event()

    async def first() -> None:
        async with gate.slot(ENDPOINT, "chat"):
            order.append("first in")
            first_in.set()
            await let_go.wait()
        order.append("first out")

    async def second() -> None:
        await first_in.wait()
        async with gate.slot(ENDPOINT, "chat"):
            order.append("second in")

    task_one = asyncio.create_task(first())
    task_two = asyncio.create_task(second())
    await first_in.wait()
    await asyncio.sleep(0)
    assert order == ["first in"], "the second call did not wait"
    let_go.set()
    await asyncio.gather(task_one, task_two)
    assert order == ["first in", "first out", "second in"]


async def test_waiting_longer_than_allowed_is_reported_as_busy(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A caller that cannot be served soon is told so, rather than being held
    until the HTTP timeout makes it look like a broken endpoint."""
    _limits(monkeypatch, parallel=1, wait=0.05)
    let_go = asyncio.Event()

    async def holder() -> None:
        async with gate.slot(ENDPOINT, "chat"):
            await let_go.wait()

    held = asyncio.create_task(holder())
    await asyncio.sleep(0)

    with pytest.raises(LLMError) as raised:
        async with gate.slot(ENDPOINT, "chat"):
            pass
    assert raised.value.code == "llm_busy"

    let_go.set()
    await held


async def test_beyond_the_queue_limit_the_answer_is_immediate(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Once far more work has arrived than the endpoint can absorb, the useful
    reply is "busy" now — not "busy" in twenty seconds."""
    _limits(monkeypatch, parallel=1, wait=5.0, queued=2)
    let_go = asyncio.Event()

    async def occupy() -> None:
        async with gate.slot(ENDPOINT, "chat"):
            await let_go.wait()

    async def queue_up() -> None:
        async with gate.slot(ENDPOINT, "chat"):
            pass

    holder = asyncio.create_task(occupy())
    await asyncio.sleep(0)
    waiters = [asyncio.create_task(queue_up()) for _ in range(2)]
    await asyncio.sleep(0)

    started = asyncio.get_running_loop().time()
    with pytest.raises(LLMError) as raised:
        async with gate.slot(ENDPOINT, "chat"):
            pass
    assert raised.value.code == "llm_busy"
    assert asyncio.get_running_loop().time() - started < 1.0

    let_go.set()
    await asyncio.gather(holder, *waiters)


async def test_endpoints_do_not_share_a_limit(monkeypatch: pytest.MonkeyPatch) -> None:
    """The chat and embedding roles usually run on different servers; a busy
    chat endpoint must not stop retrieval from embedding a query."""
    _limits(monkeypatch, parallel=1)
    let_go = asyncio.Event()

    async def occupy() -> None:
        async with gate.slot(ENDPOINT, "chat"):
            await let_go.wait()

    holder = asyncio.create_task(occupy())
    await asyncio.sleep(0)
    async with gate.slot("http://other.test/v1", "embedding"):
        pass  # reached its own slot while the first endpoint is full
    let_go.set()
    await holder


async def test_a_failed_call_gives_its_slot_back(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _limits(monkeypatch, parallel=1, wait=0.05)
    with pytest.raises(RuntimeError):
        async with gate.slot(ENDPOINT, "chat"):
            raise RuntimeError("endpoint blew up")
    # The slot is free again, so the next caller is served rather than queued.
    async with gate.slot(ENDPOINT, "chat"):
        pass


async def test_the_mode_can_tell_whether_a_turn_will_have_to_wait(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _limits(monkeypatch, parallel=1)
    assert gate.endpoint_busy(ENDPOINT) is False
    let_go = asyncio.Event()

    async def occupy() -> None:
        async with gate.slot(ENDPOINT, "chat"):
            await let_go.wait()

    holder = asyncio.create_task(occupy())
    await asyncio.sleep(0)
    assert gate.endpoint_busy(ENDPOINT) is True
    let_go.set()
    await holder
    assert gate.endpoint_busy(ENDPOINT) is False