diff --git a/sanguo_api/ws.py b/sanguo_api/ws.py new file mode 100644 index 0000000..8c96423 --- /dev/null +++ b/sanguo_api/ws.py @@ -0,0 +1,28 @@ +# sanguo_api/ws.py +"""WS 连接池:task_id → 订阅者集合。简单广播,不做重连/心跳。""" +from fastapi import WebSocket + + +class ConnectionManager: + def __init__(self): + self._connections: dict[str, set[WebSocket]] = {} + + def connect(self, task_id: str, ws: WebSocket): + self._connections.setdefault(task_id, set()).add(ws) + + def disconnect(self, task_id: str, ws: WebSocket): + conns = self._connections.get(task_id) + if conns: + conns.discard(ws) + if not conns: + del self._connections[task_id] + + async def broadcast(self, task_id: str, msg: dict): + for ws in list(self._connections.get(task_id, [])): + try: + await ws.send_json(msg) + except Exception: + self.disconnect(task_id, ws) + + +manager = ConnectionManager() diff --git a/tests/api/test_ws.py b/tests/api/test_ws.py new file mode 100644 index 0000000..6c310ff --- /dev/null +++ b/tests/api/test_ws.py @@ -0,0 +1,30 @@ +# tests/api/test_ws.py +import pytest +from unittest.mock import AsyncMock, MagicMock + + +@pytest.mark.asyncio +async def test_connect_and_broadcast(): + from sanguo_api.ws import ConnectionManager + mgr = ConnectionManager() + ws = AsyncMock() + mgr.connect("t1", ws) + assert "t1" in mgr._connections + await mgr.broadcast("t1", {"status": "running"}) + ws.send_json.assert_called_with({"status": "running"}) + + +def test_disconnect_removes_ws(): + from sanguo_api.ws import ConnectionManager + mgr = ConnectionManager() + ws = MagicMock() + mgr.connect("t1", ws) + mgr.disconnect("t1", ws) + assert ws not in mgr._connections.get("t1", set()) + + +@pytest.mark.asyncio +async def test_broadcast_no_subscribers_no_error(): + from sanguo_api.ws import ConnectionManager + mgr = ConnectionManager() + await mgr.broadcast("nope", {"x": 1}) # 不抛