Files
api/tests/test_ws_unified.py
Roberto fbd308d288 refactor(ws): rename agent_ids to scout_ids in device_hello frame
WsDeviceHello.agent_ids → scout_ids in Pydantic schema,
device_ws.py handler, and all test fixtures (test_device_ws,
test_ws_unified, test_memory_middleware). Also fixes stale
CloudAgentConfig reference in gmail.py docstring.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-05-16 01:50:15 +02:00

124 lines
4.4 KiB
Python

"""Integration tests for the unified WebSocket handler (Step 5).
Tests the device WS endpoint with home_request frames,
verifying that the correct v3 frame sequence is returned.
LLM calls are mocked to avoid network dependency.
"""
from __future__ import annotations
import json
from unittest.mock import patch
import pytest
from app.db import get_session
from app.main import app
from app.schemas import WsFrameType
from tests.conftest import TEST_USER_IDS, make_jwt
USER_ID = TEST_USER_IDS["power"]
# ── helpers ───────────────────────────────────────────────────────────────────
@pytest.fixture(autouse=True)
def _override_db(db_session):
async def _gen():
yield db_session
app.dependency_overrides[get_session] = _gen
yield
app.dependency_overrides.pop(get_session, None)
def _recv_until_end(ws, max_frames: int = 20) -> list[dict]:
"""Receive frames until stream_end or max_frames."""
frames = []
for _ in range(max_frames):
raw = ws.receive_text()
frame = json.loads(raw)
frames.append(frame)
if frame.get("type") == WsFrameType.stream_end:
break
return frames
async def _mock_home_stream(user_id, message, context):
yield "token", "Hello"
# ── tests ─────────────────────────────────────────────────────────────────────
def test_home_request_produces_stream_frames(client):
"""home_request → stream_start, stream_text+, stream_end."""
token = make_jwt("power", user_id=USER_ID)
with patch("app.api.routes.device_ws.run_home_stream", side_effect=_mock_home_stream):
with client.websocket_connect(f"/api/v1/ws/device?token={token}") as ws:
ws.send_text(json.dumps({
"type": "device_hello", "device_id": "dev-1", "scout_ids": []
}))
ws.send_text(json.dumps({
"type": "home_request",
"request_id": "r1",
"message": "List my tasks",
"conversation_history": [],
}))
frames = _recv_until_end(ws)
types = [f["type"] for f in frames]
assert WsFrameType.stream_start in types
assert WsFrameType.stream_end in types
assert types.index(WsFrameType.stream_start) < types.index(WsFrameType.stream_end)
def test_home_request_request_id_propagated(client):
"""request_id in home_request is echoed in all response frames."""
token = make_jwt("power", user_id=USER_ID)
req_id = "my-unique-req-id"
async def _stream(user_id, message, context):
yield "token", "ok"
with patch("app.api.routes.device_ws.run_home_stream", side_effect=_stream):
with client.websocket_connect(f"/api/v1/ws/device?token={token}") as ws:
ws.send_text(json.dumps({
"type": "device_hello", "device_id": "dev-3", "scout_ids": []
}))
ws.send_text(json.dumps({
"type": "home_request",
"request_id": req_id,
"message": "hello",
}))
frames = _recv_until_end(ws)
for f in frames:
if "request_id" in f:
assert f["request_id"] == req_id
def test_tool_result_dispatch_silent_on_unknown_id(client):
"""tool_result for unknown call_id is silently ignored — no crash."""
token = make_jwt("power", user_id=USER_ID)
with patch("app.api.routes.device_ws._HEARTBEAT_INTERVAL", 0.05):
with client.websocket_connect(f"/api/v1/ws/device?token={token}") as ws:
ws.send_text(json.dumps({
"type": "device_hello", "device_id": "dev-4", "scout_ids": []
}))
ws.send_text(json.dumps({
"type": "tool_result", "id": "no-such-id", "ok": True
}))
# If connection is still alive, we'll get the heartbeat ping
msg = json.loads(ws.receive_text())
assert msg["type"] == "ping"
def test_invalid_jwt_rejected(client):
"""Connection with bad token is closed before or after accept."""
with pytest.raises(Exception):
with client.websocket_connect("/api/v1/ws/device?token=badtoken") as ws:
ws.receive_text()