fastapi-sqlalchemy-asyncpg/tests/api/test_chat_ws.py

"""Tests for the websocket chat service using the local stub agent.

Uses HTTPX2's async ASGI WebSocket transport. The app lifespan is managed
explicitly because the transport does not trigger lifespan events itself.
"""

from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

import pytest
from httpx2 import AsyncClient
from httpx2.websockets import ASGIWebSocketTransport

from app.main import app

pytestmark = pytest.mark.anyio


@pytest.fixture()
async def chat_client(
    monkeypatch: pytest.MonkeyPatch,
) -> AsyncIterator[AsyncClient]:
    class RedisStub:
        async def aclose(self) -> None:
            pass

    async def get_redis_stub() -> RedisStub:
        return RedisStub()

    @asynccontextmanager
    async def scheduler_lifespan_stub() -> AsyncIterator[None]:
        yield

    monkeypatch.setattr("app.main.get_redis", get_redis_stub)
    monkeypatch.setattr("app.main.scheduler_lifespan", scheduler_lifespan_stub)

    async with app.router.lifespan_context(app):
        async with AsyncClient(
            transport=ASGIWebSocketTransport(app),
            base_url="http://testserver",
        ) as client:
            yield client


async def test_chat_ws_full_conversation_flow(chat_client: AsyncClient) -> None:
    assert await app.state.chat_sessions.count() == 0

    async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket:
        # 1. connect
        connected = await websocket.receive_json()
        assert connected["type"] == "connected"
        session_id = connected["session_id"]
        assert await app.state.chat_sessions.count() == 1

        # 2. start conversation with a system prompt
        await websocket.send_json(
            {"type": "start", "system_prompt": "You are a helpful assistant."}
        )
        started = await websocket.receive_json()
        assert started == {"type": "started", "session_id": session_id}

        # 3. send a user message and aggregate the streamed assistant chunks
        await websocket.send_json({"type": "user_message", "content": "hello"})

        event = await websocket.receive_json()
        aggregated = ""
        while event["type"] == "assistant_chunk":
            assert event["index"] >= 0
            aggregated += event["content"]
            event = await websocket.receive_json()

        assert event["type"] == "assistant_message"
        assert event["message"]["role"] == "assistant"
        assert event["message"]["content"] == aggregated
        assert "hello" in aggregated.lower()

        # 4. request full history (system + user + assistant)
        await websocket.send_json({"type": "history_request"})
        history = await websocket.receive_json()
        assert history["type"] == "history"
        roles = [m["role"] for m in history["messages"]]
        assert roles == ["system", "user", "assistant"]

        # 5. gracefully end the conversation
        await websocket.send_json({"type": "end"})

    assert await app.state.chat_sessions.count() == 0


async def test_chat_ws_invalid_event_returns_error(chat_client: AsyncClient) -> None:
    async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket:
        await websocket.receive_json()  # connected

        await websocket.send_json({"type": "not_a_real_event"})
        error = await websocket.receive_json()
        assert error["type"] == "error"

        await websocket.send_json({"type": "end"})


async def test_chat_ws_ask_question_and_print_model_answer(
    chat_client: AsyncClient,
) -> None:
    """Ask a real question over the websocket chat and print the model's answer.

    Runs against whatever ``ChatAgent`` is wired up via ``app.state.chat_agent``
    (the local stub agent by default, or a real Ollama-backed agent when
    ``CHAT_BACKEND=ollama`` is configured), exercising the full websocket
    conversation flow end-to-end.
    """
    question = "What is the capital of France?"

    async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket:
        # 1. connect
        connected = await websocket.receive_json()
        assert connected["type"] == "connected"
        session_id = connected["session_id"]

        # 2. start conversation
        await websocket.send_json(
            {"type": "start", "system_prompt": "You are a helpful assistant."}
        )
        started = await websocket.receive_json()
        assert started == {"type": "started", "session_id": session_id}

        # 3. ask the question and aggregate the streamed answer
        await websocket.send_json({"type": "user_message", "content": question})

        event = await websocket.receive_json()
        answer = ""
        while event["type"] == "assistant_chunk":
            answer += event["content"]
            event = await websocket.receive_json()

        assert event["type"] == "assistant_message"
        answer = event["message"]["content"]

        print(f"\nQuestion: {question}")
        print(f"Model answer: {answer}")

        assert isinstance(answer, str)
        assert answer.strip() != ""

        # 4. gracefully end the conversation
        await websocket.send_json({"type": "end"})