408 lines
19 KiB
Python
408 lines
19 KiB
Python
# #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 -> [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_auth_rejected_closes [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()
|
|
with patch("src.app._authenticate_websocket", return_value=False):
|
|
await websocket_endpoint(ws, "task-1")
|
|
ws.close.assert_called_once_with(code=4001, reason="Authentication required")
|
|
ws.accept.assert_not_called()
|
|
# #endregion test_auth_rejected_closes
|
|
|
|
# #region test_accepts_and_sends_initial_status [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.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_accepts_and_sends_initial_status
|
|
|
|
# #region test_source_filter [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.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_source_filter
|
|
|
|
# #region test_level_filter [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.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_level_filter
|
|
|
|
# #region test_disconnect_in_main_loop [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.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_disconnect_in_main_loop
|
|
|
|
# #region test_awaiting_input_prompt [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.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_awaiting_input_prompt
|
|
|
|
# #region test_terminal_log_triggers_close [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.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_terminal_log_triggers_close
|
|
|
|
|
|
class TestMatchesFilters:
|
|
"""matches_filters — extracted filter logic."""
|
|
|
|
# #region test_matches_filters [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_matches_filters
|
|
|
|
|
|
class TestWebSocketMainLoopCoverage:
|
|
"""Cover the main loop lines 654, 658, 660-677 — non-terminal status, log filter, log forwarding, terminal log message."""
|
|
|
|
# #region test_non_terminal_status_continue [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.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_non_terminal_status_continue
|
|
|
|
# #region test_log_filter_continue_and_forward [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)."""
|
|
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
|
|
filtered_log = _make_mock_log_entry(msg="Filtered out", source="superset_api", level="DEBUG")
|
|
# Log entry that will match
|
|
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.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-filter")
|
|
ws.accept.assert_called_once()
|
|
# #endregion test_log_filter_continue_and_forward
|
|
|
|
# #region test_terminal_log_message_detected [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.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_terminal_log_message_detected
|
|
|
|
|
|
class TestWebSocketMainLoopExceptions:
|
|
"""Generic Exception re-raise in WS main loops."""
|
|
|
|
# #region test_ws_endpoint_generic_exception [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.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_ws_endpoint_generic_exception
|
|
|
|
# #region test_task_events_generic_exception [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.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_task_events_generic_exception
|
|
|
|
# #region test_maintenance_events_generic_exception [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.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_maintenance_events_generic_exception
|
|
|
|
# #region test_dataset_ws_generic_exception [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.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_dataset_ws_generic_exception
|
|
# #endregion Test.AppModule.WsEndpoint
|