ss-tools/backend/tests/test_app_ws_endpoint.py

# #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 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