mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
Renn F
4da0245dac
F064: the bridge forwarder awaited every conn.send_text in a gather with no per-connection queue and no send timeout — one slow WS client back-pressured ALL event delivery to ALL clients (head-of-line blocking on the listen loop). Each connect_* now registers a _ClientConnection (bounded asyncio.Queue(256) + sender task); broadcasts enqueue via put_nowait (drop + structlog warn on QueueFull) and return immediately. The sender drains the queue with each send wrapped in wait_for(SEND_TIMEOUT=10s). Unregistered legacy sockets (set directly into a subscription set, bypassing connect_*) get a timeout-bounded fallback send task held in _pending_sends (ruff RUF006). disconnect cancels + drops the sender. F065: route handlers caught only WebSocketDisconnect with no finally — a non-clean exit (anyio closed-resource, CancelledError, transport error) propagated without manager.disconnect, leaking the dead socket into every subscription set forever. Added finally: manager.disconnect(websocket) to all 5 handlers (disconnect is idempotent). F066: no server-side heartbeat/idle timeout — a half-open socket from a dead container blocked receive_text forever and was never reaped. receive_text now wraps in wait_for(IDLE_TIMEOUT_SECONDS=90s); on TimeoutError, log + fall through to the F065 finally. Named module constants (no config.py precedent for WS tuning; callers/tests patch them). TDD: 22 new tests across 3 files (handler cleanup, idle timeout, send queue), non-flaky across repeats; 1 existing test adapted with a yield for the new async fan-out (assertion unchanged). ruff/mypy clean, 421 unit/api tests pass. No type:ignore/noqa.
84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
"""Operator/system WebSocket stream — ConnectionManager + the /ws/system endpoint.
|
|
|
|
The system stream carries system-wide events (rate limits) to the panel with no
|
|
per-agent keying. These tests exercise the manager's system-connection
|
|
bookkeeping and the endpoint's connect → ping/pong → disconnect lifecycle
|
|
against mock sockets (no real app/lifespan/Redis).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from fastapi import WebSocketDisconnect
|
|
from roboco.api.websocket import ConnectionManager, manager, system_stream
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_connect_system_accepts_and_tracks() -> None:
|
|
mgr = ConnectionManager()
|
|
ws = MagicMock()
|
|
ws.accept = AsyncMock()
|
|
|
|
await mgr.connect_system(ws)
|
|
|
|
ws.accept.assert_awaited_once()
|
|
assert ws in mgr.system_connections
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_broadcast_system_sends_to_every_connection() -> None:
|
|
mgr = ConnectionManager()
|
|
ws1, ws2 = MagicMock(), MagicMock()
|
|
ws1.send_text = AsyncMock()
|
|
ws2.send_text = AsyncMock()
|
|
# These sockets are placed directly into the subscription set (bypassing
|
|
# connect_system), so they take the F064 legacy fallback path: broadcast
|
|
# schedules a timeout-bounded send task per socket instead of awaiting
|
|
# send_text inline. Yield once so those tasks run before asserting.
|
|
mgr.system_connections = {ws1, ws2}
|
|
|
|
await mgr.broadcast_system({"type": "RATE_LIMIT_HIT", "provider": "anthropic"})
|
|
await asyncio.sleep(0)
|
|
|
|
ws1.send_text.assert_awaited_once()
|
|
ws2.send_text.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_broadcast_system_is_noop_when_empty() -> None:
|
|
mgr = ConnectionManager()
|
|
# No subscribers — must not raise.
|
|
await mgr.broadcast_system({"type": "x"})
|
|
|
|
|
|
def test_disconnect_removes_from_system() -> None:
|
|
mgr = ConnectionManager()
|
|
ws = MagicMock()
|
|
mgr.system_connections.add(ws)
|
|
|
|
mgr.disconnect(ws)
|
|
|
|
assert ws not in mgr.system_connections
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_system_stream_connect_ping_disconnect() -> None:
|
|
"""Endpoint sends 'connected', answers ping with pong, cleans up on close."""
|
|
ws = MagicMock()
|
|
ws.accept = AsyncMock()
|
|
ws.send_json = AsyncMock()
|
|
ws.send_text = AsyncMock()
|
|
# One ping, then the client disconnects.
|
|
ws.receive_text = AsyncMock(side_effect=["ping", WebSocketDisconnect()])
|
|
|
|
await system_stream(ws)
|
|
|
|
ws.accept.assert_awaited_once()
|
|
ws.send_json.assert_awaited_once_with({"type": "connected"})
|
|
ws.send_text.assert_awaited_with("pong")
|
|
# Disconnect handler removed the socket from the global manager.
|
|
assert ws not in manager.system_connections
|