busya/ss-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity
Files
a1cb18fad9b0e6e4a5f05a18b963cf43c9109d2d
ss-tools/backend/tests/test_app_ws_endpoint.py
busya a1cb18fad9 fix: stop WS reconnect storm on auth rejection; map 502/503/504 to NetworkError 
WebSocket endpoints now accept then close with real codes (4001 auth, 4003
permission) so clients detect auth failure via event.code instead of an opaque
403 handshake, ending the infinite reconnect storm. _authenticate_websocket
logs the actual JWT/API-key failure reason. Frontend WS consumers stop on
auth rejection and use capped exponential backoff for transient failures.

async_network.request() routes proxy 502/503/504 (HTML) responses to
NetworkError so migration/maintenance surface a clean 503 instead of a
500 JSON-parse traceback.
2026-08-01 13:23:30 +07:00

429 lines
21 KiB
Python
Raw Blame History

# #region Test.AppModule.WsEndpoint [C:3] [TYPE Module] [SEMANTICS test,app,ws,endpoint,logs]
# @BRIEF Tests for app.py — websocket_endpoint: full flow, filters, disconnect, AWAITING_INPUT, terminal log, exceptions.
# @RELATION BINDS_TO -> [App.AppModule]
import asyncio
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from starlette.websockets import WebSocketDisconnect
# ── Helpers ──────────────────────────────────────────────────────
# #region _make_mock_task [C:1] [TYPE Function]
def _make_mock_task(task_id="task-1", status="RUNNING", plugin_id="test"):
from datetime import datetime, timezone
task = MagicMock()
task.id = task_id; task.plugin_id = plugin_id; task.status = status
task.started_at = datetime(2025, 1, 15, 10, 0, 0, tzinfo=timezone.utc)
task.finished_at = None; task.user_id = "user-1"
task.result = None; task.input_required = False; task.input_request = None
task.logs = []
return task
# #endregion _make_mock_task
# #region _make_mock_log_entry [C:1] [TYPE Function]
def _make_mock_log_entry(msg="Test log", source="plugin", level="INFO"):
from datetime import datetime, timezone
ts = datetime(2025, 1, 15, 10, 0, 0, tzinfo=timezone.utc)
entry = MagicMock()
entry.source = source; entry.level = level; entry.message = msg; entry.timestamp = ts
entry.model_dump = MagicMock(return_value={
"timestamp": ts, "level": level, "message": msg, "source": source,
})
return entry
# #endregion _make_mock_log_entry
# #region _make_task_manager_mock [C:1] [TYPE Function]
def _make_task_manager_mock(task=None, logs=None):
tm = MagicMock()
tm.get_task = MagicMock(return_value=task)
tm.get_task_logs = MagicMock(return_value=logs or [])
tm.unsubscribe_logs = MagicMock()
tm.unsubscribe_status = MagicMock()
tm.unsubscribe_task_events = MagicMock()
tm.unsubscribe_maintenance_events = MagicMock()
tm.unsubscribe_dataset_events = MagicMock()
return tm
# #endregion _make_task_manager_mock
class TestWebSocketEndpointFull:
"""websocket_endpoint — full flow."""
# #region Test.AppModule.TestAuthRejectedCloses [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_auth_rejected_closes(self):
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {}; ws.close = AsyncMock(); ws.accept = AsyncMock()
with patch("src.app._authenticate_websocket", return_value=False):
await websocket_endpoint(ws, "task-1")
# Auth rejection accepts the handshake then closes with 4001 so the client can
# read event.code instead of an opaque HTTP 403 handshake rejection.
ws.accept.assert_called_once()
ws.close.assert_called_once_with(code=4001, reason="Authentication required")
# #endregion Test.AppModule.TestAuthRejectedCloses
# #region Test.AppModule.TestAcceptsAndSendsInitialStatus [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_accepts_and_sends_initial_status(self):
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(); ws.close = AsyncMock(); ws.accept = AsyncMock()
task = _make_mock_task("task-1", "RUNNING")
log = _make_mock_log_entry(msg="Starting")
lq, sq = asyncio.Queue(), asyncio.Queue()
await sq.put({"type": "task_status", "task_id": "task-1", "task": {"status": "SUCCESS"}})
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = _make_task_manager_mock(task=task, logs=[log])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
await websocket_endpoint(ws, "task-1", source=None, level=None)
ws.accept.assert_called_once()
assert ws.send_json.call_count >= 2
# #endregion Test.AppModule.TestAcceptsAndSendsInitialStatus
# #region Test.AppModule.TestSourceFilter [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_source_filter(self):
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(); ws.accept = AsyncMock()
task = _make_mock_task("task-1", "RUNNING")
plug = _make_mock_log_entry(msg="Plugin", source="plugin")
super = _make_mock_log_entry(msg="Superset", source="superset_api")
lq, sq = asyncio.Queue(), asyncio.Queue()
await sq.put({"type": "task_status", "task_id": "task-1", "task": {"status": "SUCCESS"}})
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = _make_task_manager_mock(task=task, logs=[plug, super])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
await websocket_endpoint(ws, "task-1", source="plugin")
msgs = [c[0][0].get("message") for c in ws.send_json.call_args_list if isinstance(c[0][0], dict)]
assert "Superset" not in msgs
# #endregion Test.AppModule.TestSourceFilter
# #region Test.AppModule.TestLevelFilter [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_level_filter(self):
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(); ws.accept = AsyncMock()
task = _make_mock_task("task-1", "RUNNING")
dbg = _make_mock_log_entry(msg="Debug", level="DEBUG")
inf = _make_mock_log_entry(msg="Info", level="INFO")
lq, sq = asyncio.Queue(), asyncio.Queue()
await sq.put({"type": "task_status", "task_id": "task-1", "task": {"status": "SUCCESS"}})
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = _make_task_manager_mock(task=task, logs=[dbg, inf])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
await websocket_endpoint(ws, "task-1", source=None, level="INFO")
msgs = [c[0][0].get("message") for c in ws.send_json.call_args_list if isinstance(c[0][0], dict)]
assert "Debug" not in msgs
assert "Info" in msgs
# #endregion Test.AppModule.TestLevelFilter
# #region Test.AppModule.TestDisconnectInMainLoop [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_disconnect_in_main_loop(self):
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(side_effect=[None, WebSocketDisconnect()])
ws.close = AsyncMock(); ws.accept = AsyncMock()
task = _make_mock_task("task-1", "RUNNING")
lq, sq = asyncio.Queue(), asyncio.Queue()
await sq.put({"type": "task_status", "task_id": "task-1", "task": {"status": "SUCCESS"}})
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = _make_task_manager_mock(task=task, logs=[])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
await websocket_endpoint(ws, "task-1")
ws.accept.assert_called_once()
# #endregion Test.AppModule.TestDisconnectInMainLoop
# #region Test.AppModule.TestAwaitingInputPrompt [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_awaiting_input_prompt(self):
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(); ws.accept = AsyncMock()
task = _make_mock_task("task-1", "AWAITING_INPUT")
task.input_request = {"type": "database_password", "databases": ["db1"]}
lq, sq = asyncio.Queue(), asyncio.Queue()
await sq.put({"type": "task_status", "task_id": "task-1", "task": {"status": "SUCCESS"}})
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = _make_task_manager_mock(task=task, logs=[])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
await websocket_endpoint(ws, "task-1")
msgs = [c[0][0].get("message") for c in ws.send_json.call_args_list if isinstance(c[0][0], dict)]
assert any("Task paused for user input" in (m or "") for m in msgs)
# #endregion Test.AppModule.TestAwaitingInputPrompt
# #region Test.AppModule.TestTerminalLogTriggersClose [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_terminal_log_triggers_close(self):
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(); ws.accept = AsyncMock()
task = _make_mock_task("task-1", "RUNNING")
term = _make_mock_log_entry(msg="Task completed successfully")
lq, sq = asyncio.Queue(), asyncio.Queue()
await sq.put({"type": "task_status", "task_id": "task-1", "task": {"status": "SUCCESS"}})
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = _make_task_manager_mock(task=task, logs=[term])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
await websocket_endpoint(ws, "task-1")
ws.accept.assert_called_once()
# #endregion Test.AppModule.TestTerminalLogTriggersClose
class TestMatchesFilters:
"""matches_filters — extracted filter logic."""
# #region Test.AppModule.TestMatchesFilters [C:2] [TYPE Function]
def test_matches_filters(self):
def mf(entry, source_filter=None, level_filter=None):
lh = {"DEBUG": 0, "INFO": 1, "WARNING": 2, "ERROR": 3}
ml = lh.get(level_filter, 0) if level_filter else 0
src = getattr(entry, "source", None)
if source_filter and str(src or "").lower() != source_filter:
return False
if level_filter:
el = lh.get(str(entry.level).upper(), 0)
if el < ml:
return False
return True
class E:
def __init__(self, s, l): self.source = s; self.level = l
assert mf(E("plugin", "INFO")) is True
assert mf(E("plugin", "INFO"), source_filter="other") is False
assert mf(E("other", "DEBUG"), level_filter="INFO") is False
assert mf(E("superset_api", "WARNING"), source_filter="superset_api", level_filter="WARNING") is True
assert mf(E(None, "INFO"), source_filter="plugin") is False
# #endregion Test.AppModule.TestMatchesFilters
class TestWebSocketMainLoopCoverage:
"""Cover the main loop lines 654, 658, 660-677 — non-terminal status, log filter, log forwarding, terminal log message."""
# #region Test.AppModule.TestNonTerminalStatusContinue [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_non_terminal_status_continue(self):
"""Non-terminal status (RUNNING) -> continue at line 654 is hit."""
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(); ws.accept = AsyncMock(); ws.close = AsyncMock()
task = _make_mock_task("task-continue", "RUNNING")
# Put both a non-terminal status and then a terminal one to end the loop
lq, sq = asyncio.Queue(), asyncio.Queue()
await sq.put({"type": "task_status", "task_id": "task-continue", "task": {"status": "RUNNING"}})
await sq.put({"type": "task_status", "task_id": "task-continue", "task": {"status": "SUCCESS"}})
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = _make_task_manager_mock(task=task, logs=[])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
await websocket_endpoint(ws, "task-continue")
ws.accept.assert_called_once()
# #endregion Test.AppModule.TestNonTerminalStatusContinue
# #region Test.AppModule.TestLogFilterContinueAndForward [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_log_filter_continue_and_forward(self):
"""Log entry filtered out (line 658) + log entry forwarded (lines 660-662).
Uses source filter to exercise the matches_filters rejection path.
NOTE: This test can exhibit flakiness when run in a batch due to
global task_manager singleton state leakage between test files.
"""
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(); ws.accept = AsyncMock(); ws.close = AsyncMock()
task = _make_mock_task("task-filter", "RUNNING")
# Log entry that won't match source filter — will be filtered out at line 658
filtered_log = _make_mock_log_entry(msg="Filtered out", source="superset_api", level="DEBUG")
# Log entry that will match — will be forwarded
passed_log = _make_mock_log_entry(msg="Passed through", source="plugin", level="INFO")
lq, sq = asyncio.Queue(), asyncio.Queue()
# Non-terminal status first, then the two log entries, then terminal
await sq.put({"type": "task_status", "task_id": "task-filter", "task": {"status": "RUNNING"}})
await lq.put(filtered_log)
await lq.put(passed_log)
await sq.put({"type": "task_status", "task_id": "task-filter", "task": {"status": "SUCCESS"}})
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = _make_task_manager_mock(task=task, logs=[])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
# Pass source="plugin" so superset_api log entry is filtered out (line 658)
await websocket_endpoint(ws, "task-filter", source="plugin")
ws.accept.assert_called_once()
# #endregion Test.AppModule.TestLogFilterContinueAndForward
# #region Test.AppModule.TestTerminalLogMessageDetected [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_terminal_log_message_detected(self):
"""Log entry with 'Task completed successfully' triggers delay at lines 669-677."""
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(); ws.accept = AsyncMock(); ws.close = AsyncMock()
task = _make_mock_task("task-term-msg", "RUNNING")
term_log = _make_mock_log_entry(msg="Task completed successfully", source="plugin", level="INFO")
lq, sq = asyncio.Queue(), asyncio.Queue()
await sq.put({"type": "task_status", "task_id": "task-term-msg", "task": {"status": "RUNNING"}})
await lq.put(term_log)
await sq.put({"type": "task_status", "task_id": "task-term-msg", "task": {"status": "SUCCESS"}})
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
patch("src.app.logger") as mock_logger,
):
tm = _make_task_manager_mock(task=task, logs=[])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
await websocket_endpoint(ws, "task-term-msg")
ws.accept.assert_called_once()
# #endregion Test.AppModule.TestTerminalLogMessageDetected
class TestWebSocketMainLoopExceptions:
"""Generic Exception re-raise in WS main loops."""
# #region Test.AppModule.TestWsEndpointGenericException [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_ws_endpoint_generic_exception(self):
from src.app import websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(); ws.accept = AsyncMock()
task = _make_mock_task("task-1", "RUNNING")
lq, sq = asyncio.Queue(), asyncio.Queue()
await sq.put("not-a-dict") # will cause AttributeError
async def sl(t): return lq
async def ss(t): return sq
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = _make_task_manager_mock(task=task, logs=[])
tm.subscribe_logs = sl; tm.subscribe_status = ss; mg.return_value = tm
with pytest.raises(Exception):
await websocket_endpoint(ws, "task-1")
# #endregion Test.AppModule.TestWsEndpointGenericException
# #region Test.AppModule.TestTaskEventsGenericException [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_task_events_generic_exception(self):
from src.app import task_events_websocket
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(side_effect=RuntimeError("crash"))
ws.accept = AsyncMock()
q = asyncio.Queue()
await q.put({"type": "task_status"})
async def sub(): return q
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = MagicMock(); tm.subscribe_task_events = sub; tm.unsubscribe_task_events = MagicMock()
mg.return_value = tm
with pytest.raises(RuntimeError, match="crash"):
await task_events_websocket(ws)
# #endregion Test.AppModule.TestTaskEventsGenericException
# #region Test.AppModule.TestMaintenanceEventsGenericException [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_maintenance_events_generic_exception(self):
from src.app import maintenance_events_websocket
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(side_effect=ValueError("bad"))
ws.accept = AsyncMock()
q = asyncio.Queue(); await q.put({"type": "maintenance.event_created"})
async def sub(): return q
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = MagicMock(); tm.subscribe_maintenance_events = sub; tm.unsubscribe_maintenance_events = MagicMock()
mg.return_value = tm
with pytest.raises(ValueError, match="bad"):
await maintenance_events_websocket(ws)
# #endregion Test.AppModule.TestMaintenanceEventsGenericException
# #region Test.AppModule.TestDatasetWsGenericException [C:2] [TYPE Function]
@pytest.mark.asyncio
async def test_dataset_ws_generic_exception(self):
from src.app import dataset_websocket_endpoint
ws = MagicMock(); ws.query_params = {"token": "valid"}
ws.send_json = AsyncMock(side_effect=RuntimeError("fail"))
ws.accept = AsyncMock()
q = asyncio.Queue(); await q.put({"type": "dataset.updated"})
async def sub(e): return q
with (
patch("src.app._authenticate_websocket", return_value=True),
patch("src.app._authorize_websocket", return_value=True),
patch("src.app.get_task_manager") as mg,
):
tm = MagicMock(); tm.subscribe_dataset_events = sub; tm.unsubscribe_dataset_events = MagicMock()
mg.return_value = tm
await dataset_websocket_endpoint(ws, "env-1")
ws.accept.assert_called_once()
# #endregion Test.AppModule.TestDatasetWsGenericException
# #endregion Test.AppModule.WsEndpoint
Reference in New Issue View Git Blame Copy Permalink