beaver_project/app-instance/backend/tests/unit/test_websocket_chat.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from fastapi.testclient import TestClient

from beaver.interfaces.web.app import create_app
from beaver.services.agent_service import AgentService


@dataclass(slots=True)
class StubRunResult:
    session_id: str
    run_id: str = "run-1"
    output_text: str = "ok"
    finish_reason: str = "stop"
    tool_iterations: int = 0
    provider_name: str | None = "stub"
    model: str | None = "stub-model"
    usage: dict[str, Any] = field(default_factory=lambda: {"total_tokens": 3})
    task_id: str | None = "task-1"
    task_status: str | None = "awaiting_feedback"
    validation_result: dict[str, Any] | None = field(default_factory=lambda: {"accepted": True})


class StubAgentService(AgentService):
    def __init__(self, *, fail: bool = False) -> None:
        super().__init__()
        self.fail = fail
        self.calls: list[dict[str, Any]] = []

    async def submit_direct(self, message: str, **kwargs: Any) -> StubRunResult:  # type: ignore[override]
        self.calls.append({"message": message, **kwargs})
        if self.fail:
            raise RuntimeError("boom")
        return StubRunResult(
            session_id=kwargs.get("session_id") or "web:default",
            output_text=f"echo:{message}",
        )


def test_websocket_ping_pong() -> None:
    app = create_app(service=StubAgentService(), manage_service_lifecycle=False)

    with TestClient(app) as client:
        with client.websocket_connect("/ws/web:alpha") as websocket:
            websocket.send_json({"type": "ping"})
            assert websocket.receive_json() == {"type": "pong"}


def test_websocket_message_returns_chat_metadata_and_session_updated() -> None:
    service = StubAgentService()
    app = create_app(service=service, manage_service_lifecycle=False)

    with TestClient(app) as client:
        with client.websocket_connect("/ws/web:alpha") as websocket:
            websocket.send_json(
                {
                    "type": "message",
                    "content": "hello",
                    "metadata": {"source": "test"},
                    "attachments": [{"file_id": "file-1", "name": "a.txt"}],
                }
            )
            assert websocket.receive_json() == {"type": "status", "status": "thinking"}
            message = websocket.receive_json()
            session_updated = websocket.receive_json()

    assert service.calls == [
        {
            "message": "hello",
            "session_id": "web:alpha",
            "source": "websocket",
            "user_id": None,
            "title": None,
            "execution_context": None,
            "model": None,
            "provider_name": None,
            "embedding_model": None,
        }
    ]
    assert message["type"] == "message"
    assert message["role"] == "assistant"
    assert message["content"] == "echo:hello"
    assert message["session_id"] == "web:alpha"
    assert message["run_id"] == "run-1"
    assert message["task_id"] == "task-1"
    assert message["task_status"] == "awaiting_feedback"
    assert message["validation_result"] == {"accepted": True}
    assert message["validation_status"] == "passed"
    assert message["metadata"]["input_metadata"] == {
        "source": "test",
        "attachments": [{"file_id": "file-1", "name": "a.txt"}],
    }
    assert session_updated == {
        "type": "session_updated",
        "session_id": "web:alpha",
        "source": "websocket",
    }


def test_websocket_empty_content_returns_error_without_runtime_call() -> None:
    service = StubAgentService()
    app = create_app(service=service, manage_service_lifecycle=False)

    with TestClient(app) as client:
        with client.websocket_connect("/ws/web:alpha") as websocket:
            websocket.send_json({"type": "message", "content": "   "})
            assert websocket.receive_json() == {"type": "error", "error": "'content' is required"}

    assert service.calls == []


def test_websocket_runtime_error_returns_assistant_error_message() -> None:
    service = StubAgentService(fail=True)
    app = create_app(service=service, manage_service_lifecycle=False)

    with TestClient(app) as client:
        with client.websocket_connect("/ws/web:alpha") as websocket:
            websocket.send_json({"type": "message", "content": "hello"})
            assert websocket.receive_json() == {"type": "status", "status": "thinking"}
            message = websocket.receive_json()
            websocket.send_json({"type": "ping"})
            pong = websocket.receive_json()

    assert message["type"] == "message"
    assert message["role"] == "assistant"
    assert message["session_id"] == "web:alpha"
    assert message["finish_reason"] == "error"
    assert "boom" in message["content"]
    assert pong == {"type": "pong"}