brain-of-reese/tests/e2e/test_chat_db_pool.py

"""E2E: chat concurrency cap (SEC-14-04, phase 106, task 03).

Verifies the concurrency cap end-to-end using the app server:
- The app handles concurrent chat requests correctly.
- Requests exceeding the concurrency cap get a 503 response.

Requires:  podman compose up -d db
"""
from __future__ import annotations

import math
import re
from collections.abc import AsyncIterator
from typing import Any

import pytest
from httpx import ASGITransport, AsyncClient

from app.api import chat as chat_api
from app.main import app as fastapi_app
from app.rag.llm import StreamPiece, ToolCallPiece

# The tests/conftest.py sets BOR_ADMIN_PASSWORD="test-admin-password"
# before importing app.main — use the same value here.
ADMIN_PASSWORD = "test-admin-password"

_DIM = 768
_TOKEN_RE = re.compile(r"[a-z0-9]+")


def _token_vec(text: str) -> list[float]:
    """Bag-of-words unit vector — same algorithm as the E2E mock."""
    import hashlib

    vec = [0.0] * _DIM
    for tok in _TOKEN_RE.findall(text.lower()):
        vec[int(hashlib.md5(tok.encode()).hexdigest(), 16) % _DIM] += 1.0
    norm = math.sqrt(sum(v * v for v in vec)) or 1.0
    return [v / norm for v in vec]


class FakeChatLLM:
    """Minimal duck-typed LLM client for the chat endpoint."""

    def __init__(self, answer: str = "Test answer.") -> None:
        self.answer = answer

    async def embed_one(self, message: str) -> list[float]:
        return _token_vec(message)

    async def chat_stream(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        scaffolding: Any = None,
    ) -> AsyncIterator[StreamPiece | ToolCallPiece]:
        yield StreamPiece("thinking", "thinking")
        yield StreamPiece("content", self.answer)


def _mock_llm_fixture(app) -> FakeChatLLM:
    """Set up a fake LLM on the app and return it."""
    fake = FakeChatLLM(answer="Test answer.")
    app.dependency_overrides[chat_api.get_llm] = lambda: fake
    return fake


class TestConcurrencyCapE2E:
    """E2E tests for the chat concurrency cap using the real app."""

    @pytest.mark.anyio
    async def test_chat_within_cap_works(self) -> None:
        """A single chat request within the cap succeeds (returns stream)."""
        _mock_llm_fixture(fastapi_app)
        transport = ASGITransport(app=fastapi_app)
        async with AsyncClient(transport=transport, base_url="http://test") as client:
            # Login first
            resp = await client.post(
                "/api/login",
                json={"password": ADMIN_PASSWORD},
            )
            assert resp.status_code == 204, f"Login failed: {resp.status_code}"

            # Send a chat request — should get a streaming response
            resp = await client.post(
                "/api/chat",
                json={"message": "hello"},
            )
            assert resp.status_code == 200
            assert resp.headers["content-type"].startswith("text/event-stream")

    @pytest.mark.anyio
    async def test_exceeding_cap_gets_503(self) -> None:
        """Requests exceeding the concurrency cap get 503."""
        import app.api.chat as chat_module

        # Reset state
        chat_module._chat_active = 0
        chat_module._chat_semaphore = None

        _mock_llm_fixture(fastapi_app)

        transport = ASGITransport(app=fastapi_app)
        async with AsyncClient(transport=transport, base_url="http://test") as client:
            # Login first
            resp = await client.post(
                "/api/login",
                json={"password": ADMIN_PASSWORD},
            )
            assert resp.status_code == 204

            # Fill up the slots
            chat_module._chat_active = 10  # default cap

            try:
                # This request should get 503
                resp = await client.post(
                    "/api/chat",
                    json={"message": "overflow"},
                )
                assert resp.status_code == 503
                body = resp.json()
                assert "Too many concurrent" in body["detail"]
            finally:
                chat_module._chat_active = 0

    @pytest.mark.anyio
    async def test_slot_released_after_stream(self) -> None:
        """After a stream completes, the slot is freed for the next request."""
        import app.api.chat as chat_module

        # Reset state
        chat_module._chat_active = 0
        chat_module._chat_semaphore = None

        _mock_llm_fixture(fastapi_app)

        transport = ASGITransport(app=fastapi_app)
        async with AsyncClient(transport=transport, base_url="http://test") as client:
            # Login first
            resp = await client.post(
                "/api/login",
                json={"password": ADMIN_PASSWORD},
            )
            assert resp.status_code == 204

            # Simulate a stream completing (counter back to 0)
            chat_module._chat_active = 0

            # Request should succeed
            resp = await client.post(
                "/api/chat",
                json={"message": "after release"},
            )
            assert resp.status_code == 200