feat(api): WS 连接池 ConnectionManager

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()

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})  # 不抛