53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
"""WebSocket 快速重连防回归测试。"""
|
|
import asyncio
|
|
import sys
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "app"))
|
|
|
|
from services.ws_hub import WebSocketHub
|
|
|
|
|
|
class FakeWebSocket:
|
|
def __init__(self):
|
|
self.accept = AsyncMock()
|
|
self.send_json = AsyncMock()
|
|
|
|
|
|
def test_old_socket_disconnect_does_not_remove_reconnected_socket():
|
|
async def scenario():
|
|
hub = WebSocketHub()
|
|
old_ws = FakeWebSocket()
|
|
new_ws = FakeWebSocket()
|
|
await hub.connect(old_ws, "device-1")
|
|
hub.device_info["device-1"] = {"project_id": "p1", "model": "phone"}
|
|
await hub.connect(new_ws, "device-1")
|
|
|
|
with patch("services.ws_hub._record_device_event", new=AsyncMock()) as record:
|
|
await hub.disconnect("device-1", websocket=old_ws)
|
|
|
|
assert hub.connections["device-1"] is new_ws
|
|
assert "device-1" in hub.device_info
|
|
record.assert_not_awaited()
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_disconnect_is_idempotent_and_records_once():
|
|
async def scenario():
|
|
hub = WebSocketHub()
|
|
ws = FakeWebSocket()
|
|
await hub.connect(ws, "device-1")
|
|
hub.device_info["device-1"] = {"project_id": "p1", "model": "phone"}
|
|
|
|
with patch("services.ws_hub._record_device_event", new=AsyncMock()) as record:
|
|
await hub.disconnect("device-1", websocket=ws)
|
|
await hub.disconnect("device-1", websocket=ws)
|
|
|
|
assert "device-1" not in hub.connections
|
|
assert "device-1" not in hub.device_info
|
|
record.assert_awaited_once()
|
|
|
|
asyncio.run(scenario())
|