"""Tests for WebSocket status endpoint and broadcast helpers.""" import json import pytest from fastapi.testclient import TestClient from starlette.websockets import WebSocketDisconnect from app.api.routes.status import _connections, broadcast_scan_update, broadcast_status from app.main import app # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_token() -> str: from app.core.security import create_access_token return create_access_token("admin") # --------------------------------------------------------------------------- # WebSocket authentication # --------------------------------------------------------------------------- def test_websocket_rejected_without_token(): """Connection that sends no token field must be closed with 1008.""" with TestClient(app) as client, pytest.raises(WebSocketDisconnect), client.websocket_connect("/api/v1/status/ws/status") as ws: ws.send_text(json.dumps({})) # missing token field ws.receive_text() # triggers WebSocketDisconnect from server close def test_websocket_rejected_with_invalid_token(): """Connection that sends a garbage token must be closed.""" with TestClient(app) as client, pytest.raises(WebSocketDisconnect), client.websocket_connect("/api/v1/status/ws/status") as ws: ws.send_text(json.dumps({"token": "not-a-valid-jwt"})) ws.receive_text() def test_websocket_rejected_with_malformed_json(): """Connection that sends non-JSON as auth must be closed.""" with TestClient(app) as client, pytest.raises(WebSocketDisconnect), client.websocket_connect("/api/v1/status/ws/status") as ws: ws.send_text("not-json") ws.receive_text() def test_websocket_accepted_with_valid_token(): """Connection that sends a valid JWT as first message must be accepted.""" token = _make_token() with TestClient(app) as client, client.websocket_connect("/api/v1/status/ws/status") as ws: ws.send_text(json.dumps({"token": token})) # Connection is open — subsequent messages should not raise ws.send_text("ping") # --------------------------------------------------------------------------- # broadcast_status # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_broadcast_status_sends_to_connected_clients(): """broadcast_status sends a JSON message to all active connections.""" received: list[str] = [] class FakeWS: async def send_text(self, text: str) -> None: received.append(text) fake = FakeWS() _connections.append(fake) try: await broadcast_status( node_id="node-1", status="online", checked_at="2024-01-01T00:00:00", response_time_ms=42, ) finally: _connections.remove(fake) assert len(received) == 1 msg = json.loads(received[0]) assert msg["type"] == "status" assert msg["node_id"] == "node-1" assert msg["status"] == "online" assert msg["response_time_ms"] == 42 @pytest.mark.asyncio async def test_broadcast_status_no_response_time(): """response_time_ms defaults to None.""" received: list[str] = [] class FakeWS: async def send_text(self, text: str) -> None: received.append(text) fake = FakeWS() _connections.append(fake) try: await broadcast_status(node_id="n", status="offline", checked_at="t") finally: _connections.remove(fake) msg = json.loads(received[0]) assert msg["response_time_ms"] is None @pytest.mark.asyncio async def test_broadcast_status_removes_dead_connection(): """A connection that raises on send is removed from _connections.""" class DeadWS: async def send_text(self, _: str) -> None: raise RuntimeError("disconnected") dead = DeadWS() _connections.append(dead) initial_len = len(_connections) await broadcast_status(node_id="n", status="online", checked_at="t") assert dead not in _connections assert len(_connections) == initial_len - 1 # --------------------------------------------------------------------------- # broadcast_scan_update # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_broadcast_scan_update(): """broadcast_scan_update sends scan_device_found payload.""" received: list[str] = [] class FakeWS: async def send_text(self, text: str) -> None: received.append(text) fake = FakeWS() _connections.append(fake) try: await broadcast_scan_update(run_id="run-42", devices_found=3) finally: _connections.remove(fake) assert len(received) == 1 msg = json.loads(received[0]) assert msg["type"] == "scan_device_found" assert msg["run_id"] == "run-42" assert msg["devices_found"] == 3 @pytest.mark.asyncio async def test_broadcast_no_connections(): """broadcast_* with no connections must not raise.""" assert len(_connections) == 0 await broadcast_status(node_id="n", status="online", checked_at="t") await broadcast_scan_update(run_id="r", devices_found=0)