diff --git a/shelfmark/core/socket_handlers.py b/shelfmark/core/socket_handlers.py new file mode 100644 index 00000000..d5cc3c67 --- /dev/null +++ b/shelfmark/core/socket_handlers.py @@ -0,0 +1,106 @@ +"""Socket.IO handlers: connect, disconnect and a manual status request.""" + +from __future__ import annotations + +import sqlite3 +from typing import TYPE_CHECKING + +from flask import request +from flask_socketio import emit + +from shelfmark.core.logger import setup_logger +from shelfmark.core.viewer_redaction import redact_status +from shelfmark.download import orchestrator as backend + +if TYPE_CHECKING: + from collections.abc import Callable + + from flask_socketio import SocketIO + + from shelfmark.api.websocket import WebSocketManager + +logger = setup_logger(__name__) + +_OPERATIONAL_ERRORS = (OSError, RuntimeError, TypeError, ValueError, sqlite3.Error) + + +def _get_request_sid() -> str | None: + """Return the Socket.IO session id for the active request when available.""" + sid = getattr(request, "sid", None) + return sid if isinstance(sid, str) and sid else None + + +def register_socket_handlers( + socketio: SocketIO, + *, + ws_manager: WebSocketManager, + resolve_status_scope: Callable[..., tuple[bool, int | None, bool]], +) -> None: + """Register the connect, disconnect and request_status handlers. + + ``resolve_status_scope`` stays in ``main.py`` because the HTTP routes use it too. + """ + + @socketio.on("connect") + def handle_connect() -> None: + """Handle client connection.""" + logger.info("WebSocket client connected") + + # Track the connection (triggers warmup callbacks on first connect) + ws_manager.client_connected() + + # Join appropriate room based on authenticated user session + is_admin, db_user_id, can_access_status = resolve_status_scope() + sid = _get_request_sid() + if sid is None: + logger.warning("Socket.IO connect event missing sid") + return + ws_manager.join_user_room(sid, is_admin=is_admin, db_user_id=db_user_id) + + # Send initial status to the newly connected client (filtered) + try: + if not can_access_status: + emit("status_update", {}) + return + + user_id = None if is_admin else db_user_id + status = backend.queue_status(user_id=user_id) + emit("status_update", status if is_admin else redact_status(status)) + except _OPERATIONAL_ERRORS: + logger.exception("Error sending initial status") + + @socketio.on("disconnect") + def handle_disconnect() -> None: + """Handle client disconnection.""" + logger.info("WebSocket client disconnected") + + # Leave room + sid = _get_request_sid() + if sid is not None: + ws_manager.leave_user_room(sid) + + # Track the disconnection + ws_manager.client_disconnected() + + @socketio.on("request_status") + def handle_status_request() -> None: + """Handle manual status request from client.""" + try: + is_admin, db_user_id, can_access_status = resolve_status_scope() + sid = _get_request_sid() + if sid is None: + logger.warning("Socket.IO request_status event missing sid") + emit("status_update", {}) + return + ws_manager.sync_user_room(sid, is_admin=is_admin, db_user_id=db_user_id) + + if not can_access_status: + emit("status_update", {}) + return + + user_id = None if is_admin else db_user_id + status = backend.queue_status(user_id=user_id) + emit("status_update", status if is_admin else redact_status(status)) + except _OPERATIONAL_ERRORS: + logger.exception("Error handling status request") + emit("error", {"message": "Failed to get status"}) diff --git a/shelfmark/main.py b/shelfmark/main.py index 50cb09be..1f07da61 100644 --- a/shelfmark/main.py +++ b/shelfmark/main.py @@ -15,7 +15,7 @@ from urllib.parse import quote from flask import Flask, g, jsonify, request, send_file, session from flask_cors import CORS -from flask_socketio import SocketIO, emit +from flask_socketio import SocketIO from werkzeug.middleware.proxy_fix import ProxyFix from werkzeug.security import check_password_hash from werkzeug.wrappers import Response @@ -102,6 +102,7 @@ from shelfmark.core.requests_service import ( sync_delivery_states_from_queue_status, ) from shelfmark.core.settings_routes import register_settings_routes +from shelfmark.core.socket_handlers import register_socket_handlers from shelfmark.core.socketio_origins import build_cors_allowed_origins from shelfmark.core.stats_routes import STATS_PATH, register_stats_routes from shelfmark.core.system_routes import DebugHooks, register_system_routes @@ -2471,78 +2472,9 @@ register_settings_routes(app, login_required) register_onboarding_routes(app, login_required) -def _get_request_sid() -> str | None: - """Return the Socket.IO session id for the active request when available.""" - sid = getattr(request, "sid", None) - return sid if isinstance(sid, str) and sid else None - - -# WebSocket event handlers -@socketio.on("connect") -def handle_connect() -> None: - """Handle client connection.""" - logger.info("WebSocket client connected") - - # Track the connection (triggers warmup callbacks on first connect) - ws_manager.client_connected() - - # Join appropriate room based on authenticated user session - is_admin, db_user_id, can_access_status = _resolve_status_scope() - sid = _get_request_sid() - if sid is None: - logger.warning("Socket.IO connect event missing sid") - return - ws_manager.join_user_room(sid, is_admin=is_admin, db_user_id=db_user_id) - - # Send initial status to the newly connected client (filtered) - try: - if not can_access_status: - emit("status_update", {}) - return - - user_id = None if is_admin else db_user_id - status = backend.queue_status(user_id=user_id) - emit("status_update", status if is_admin else redact_status(status)) - except _OPERATIONAL_ERRORS: - logger.exception("Error sending initial status") - - -@socketio.on("disconnect") -def handle_disconnect() -> None: - """Handle client disconnection.""" - logger.info("WebSocket client disconnected") - - # Leave room - sid = _get_request_sid() - if sid is not None: - ws_manager.leave_user_room(sid) - - # Track the disconnection - ws_manager.client_disconnected() - - -@socketio.on("request_status") -def handle_status_request() -> None: - """Handle manual status request from client.""" - try: - is_admin, db_user_id, can_access_status = _resolve_status_scope() - sid = _get_request_sid() - if sid is None: - logger.warning("Socket.IO request_status event missing sid") - emit("status_update", {}) - return - ws_manager.sync_user_room(sid, is_admin=is_admin, db_user_id=db_user_id) - - if not can_access_status: - emit("status_update", {}) - return - - user_id = None if is_admin else db_user_id - status = backend.queue_status(user_id=user_id) - emit("status_update", status if is_admin else redact_status(status)) - except _OPERATIONAL_ERRORS: - logger.exception("Error handling status request") - emit("error", {"message": "Failed to get status"}) +register_socket_handlers( + socketio, ws_manager=ws_manager, resolve_status_scope=_resolve_status_scope +) logger.log_resource_usage() diff --git a/tests/core/test_socket_handlers.py b/tests/core/test_socket_handlers.py new file mode 100644 index 00000000..e3caa1de --- /dev/null +++ b/tests/core/test_socket_handlers.py @@ -0,0 +1,181 @@ +"""Characterization tests for the Socket.IO handlers, written before they moved out of main.py. + +They pin what a client sees on connect, on a manual status request and on disconnect, and how +the room bookkeeping is called, for no auth, an admin, a plain user and an anonymous visitor. +""" + +from __future__ import annotations + +import importlib +from unittest.mock import patch + +import pytest + + +@pytest.fixture(scope="module") +def main_module(): + """Import `shelfmark.main` with background startup disabled.""" + with patch("shelfmark.download.orchestrator.start"): + import shelfmark.main as main + + importlib.reload(main) + return main + + +@pytest.fixture +def client(main_module): + return main_module.app.test_client() + + +def _login(client, *, user_id="alice", db_user_id: int | None = 7, is_admin=False) -> None: + with client.session_transaction() as sess: + sess["user_id"] = user_id + sess["is_admin"] = is_admin + if db_user_id is not None: + sess["db_user_id"] = db_user_id + + +def _socket(main_module, client): + return main_module.socketio.test_client(main_module.app, flask_test_client=client) + + +def _events(socket) -> list[tuple[str, dict]]: + return [(item["name"], item["args"][0]) for item in socket.get_received()] + + +_STATUS = {"queued": {"task-1": {"title": "A Book", "source": "direct_download"}}} + + +class TestConnect: + def test_no_auth_joins_as_admin_and_gets_the_whole_status(self, main_module, client): + with patch.object(main_module, "get_auth_mode", return_value="none"): + with patch.object(main_module.backend, "queue_status", return_value=_STATUS) as status: + with patch.object(main_module.ws_manager, "join_user_room") as join: + with patch.object(main_module.ws_manager, "client_connected") as connected: + socket = _socket(main_module, client) + events = _events(socket) + socket.disconnect() + + assert events == [("status_update", _STATUS)] + status.assert_called_once_with(user_id=None) + connected.assert_called_once_with() + assert join.call_args.kwargs == {"is_admin": True, "db_user_id": None} + + def test_an_admin_gets_the_whole_status(self, main_module, client): + _login(client, user_id="root", db_user_id=None, is_admin=True) + with patch.object(main_module, "get_auth_mode", return_value="builtin"): + with patch.object(main_module.backend, "queue_status", return_value=_STATUS): + with patch.object(main_module.ws_manager, "join_user_room") as join: + socket = _socket(main_module, client) + events = _events(socket) + socket.disconnect() + + assert events == [("status_update", _STATUS)] + assert join.call_args.kwargs == {"is_admin": True, "db_user_id": None} + + def test_a_user_gets_their_own_status_with_release_details_hidden(self, main_module, client): + _login(client, user_id="alice", db_user_id=7) + with patch.object(main_module, "get_auth_mode", return_value="builtin"): + with patch.object(main_module.backend, "queue_status", return_value=_STATUS) as status: + with patch( + "shelfmark.core.socket_handlers.redact_status", return_value={"redacted": {}} + ): + with patch.object(main_module.ws_manager, "join_user_room") as join: + socket = _socket(main_module, client) + events = _events(socket) + socket.disconnect() + + assert events == [("status_update", {"redacted": {}})] + status.assert_called_once_with(user_id=7) + assert join.call_args.kwargs == {"is_admin": False, "db_user_id": 7} + + def test_a_visitor_with_no_session_gets_an_empty_status(self, main_module, client): + with patch.object(main_module, "get_auth_mode", return_value="builtin"): + with patch.object(main_module.backend, "queue_status") as status: + with patch.object(main_module.ws_manager, "join_user_room"): + socket = _socket(main_module, client) + events = _events(socket) + socket.disconnect() + + assert events == [("status_update", {})] + status.assert_not_called() + + def test_a_failure_reading_the_status_sends_nothing_and_keeps_the_connection( + self, main_module, client + ): + _login(client, user_id="root", db_user_id=None, is_admin=True) + with patch.object(main_module, "get_auth_mode", return_value="builtin"): + with patch.object( + main_module.backend, "queue_status", side_effect=RuntimeError("broke") + ): + with patch.object(main_module.ws_manager, "join_user_room"): + socket = _socket(main_module, client) + events = _events(socket) + connected = socket.is_connected() + socket.disconnect() + + assert events == [] + assert connected is True + + +class TestRequestStatus: + def test_a_manual_request_syncs_the_room_and_answers_with_the_status(self, main_module, client): + _login(client, user_id="root", db_user_id=None, is_admin=True) + with patch.object(main_module, "get_auth_mode", return_value="builtin"): + with patch.object(main_module.backend, "queue_status", return_value=_STATUS): + with patch.object(main_module.ws_manager, "join_user_room"): + socket = _socket(main_module, client) + socket.get_received() + with patch.object(main_module.ws_manager, "sync_user_room") as sync: + socket.emit("request_status") + events = _events(socket) + socket.disconnect() + + assert events == [("status_update", _STATUS)] + assert sync.call_args.kwargs == {"is_admin": True, "db_user_id": None} + + def test_a_visitor_with_no_session_gets_an_empty_status(self, main_module, client): + with patch.object(main_module, "get_auth_mode", return_value="builtin"): + with patch.object(main_module.ws_manager, "join_user_room"): + socket = _socket(main_module, client) + socket.get_received() + with patch.object(main_module.ws_manager, "sync_user_room"): + socket.emit("request_status") + events = _events(socket) + socket.disconnect() + + assert events == [("status_update", {})] + + def test_a_failure_answers_with_an_error_event(self, main_module, client): + _login(client, user_id="root", db_user_id=None, is_admin=True) + with patch.object(main_module, "get_auth_mode", return_value="builtin"): + with patch.object(main_module.backend, "queue_status", return_value=_STATUS): + with patch.object(main_module.ws_manager, "join_user_room"): + socket = _socket(main_module, client) + socket.get_received() + with patch.object( + main_module.backend, "queue_status", side_effect=RuntimeError("broke") + ): + with patch.object(main_module.ws_manager, "sync_user_room"): + socket.emit("request_status") + events = _events(socket) + socket.disconnect() + + assert events == [("error", {"message": "Failed to get status"})] + + +class TestDisconnect: + def test_leaves_the_room_and_counts_the_disconnect(self, main_module, client): + with patch.object(main_module, "get_auth_mode", return_value="none"): + with patch.object(main_module.backend, "queue_status", return_value={}): + with patch.object(main_module.ws_manager, "join_user_room"): + socket = _socket(main_module, client) + with patch.object(main_module.ws_manager, "leave_user_room") as leave: + with patch.object( + main_module.ws_manager, "client_disconnected" + ) as disconnected: + socket.disconnect() + + leave.assert_called_once() + assert isinstance(leave.call_args.args[0], str) and leave.call_args.args[0] + disconnected.assert_called_once_with()