From 0b6319194f5dbe5e08c1405a9e0ea60adba0f7cb Mon Sep 17 00:00:00 2001 From: Fabian Zills Date: Sun, 18 Jan 2026 13:39:21 +0100 Subject: [PATCH] feat: handle socket reconnection after server restart Add connect_error handler in SocketManager that re-registers users when the server rejects connection with "User not found" error. This happens when the server restarts and Redis data is cleared - the client's JWT token is still valid but the user no longer exists in Redis. The handler: - Detects "User not found" errors during reconnection - Re-calls api.login() to re-register the user in Redis - Updates the JWT token so subsequent reconnect attempts succeed Co-Authored-By: Claude Opus 4.5 --- src/zndraw/socket_manager.py | 37 ++++++ tests/test_socket_reconnection.py | 183 ++++++++++++++++++++++++++++++ 2 files changed, 220 insertions(+) create mode 100644 tests/test_socket_reconnection.py diff --git a/src/zndraw/socket_manager.py b/src/zndraw/socket_manager.py index 9b15254b4..566b794d4 100644 --- a/src/zndraw/socket_manager.py +++ b/src/zndraw/socket_manager.py @@ -22,6 +22,7 @@ def __init__(self, zndraw_instance: "ZnDraw"): def _register_handlers(self): self.sio.on("connect", self._on_connect) + self.sio.on("connect_error", self._on_connect_error) self.sio.on(SocketEvents.FRAME_UPDATE, self._on_frame_update) self.sio.on(SocketEvents.SELECTION_UPDATE, self._on_selection_update) self.sio.on(SocketEvents.ROOM_UPDATE, self._on_room_update) @@ -97,6 +98,42 @@ def connect(self): # Mark initial connection as done - subsequent connects are reconnects self._initial_connect_done = True + def _on_connect_error(self, data): + """Handle socket connection error. + + When the server rejects connection with "User not found" (e.g., after + server restart when Redis data is cleared), re-register the user + via the login API and update the JWT token for the next reconnect attempt. + """ + error_msg = str(data) if data else "" + log.warning(f"Socket connection error: {error_msg}") + + # Check if error is "User not found" - need to re-register + if "User not found" in error_msg or "not found" in error_msg.lower(): + log.info("User not found on server, re-registering...") + try: + # Re-call login API to re-register the user + # This creates the user in Redis and gets a fresh JWT token + login_data = self.zndraw.api.login( + user_name=self.zndraw.user, + password=( + self.zndraw.password.get_secret_value() + if self.zndraw.password + else None + ), + ) + + # Update user info from response (in case it changed) + self.zndraw.user = login_data["userName"] + self.zndraw.role = login_data.get("role", "guest") + + log.info( + f"Re-registered as {self.zndraw.user}, " + "next reconnect attempt should succeed" + ) + except Exception as e: + log.error(f"Failed to re-register user: {e}") + def disconnect(self): if self.sio.connected: self.sio.disconnect() diff --git a/tests/test_socket_reconnection.py b/tests/test_socket_reconnection.py new file mode 100644 index 000000000..387584693 --- /dev/null +++ b/tests/test_socket_reconnection.py @@ -0,0 +1,183 @@ +"""Tests for socket reconnection handling after server restart. + +Tests the _on_connect_error handler in socket_manager.py that handles +the case where a user is not found after server restart. +""" + +import uuid +from unittest.mock import MagicMock, patch + +import pytest + +from zndraw import ZnDraw + + +def test_connect_error_handler_reregisters_user_on_not_found(server): + """Test that _on_connect_error re-registers user when 'User not found' error occurs.""" + # Create a client and connect normally + room = uuid.uuid4().hex + client = ZnDraw(room=room, url=server) + assert client.socket.connected + + # Store original user info + original_user = client.user + original_token = client.api.jwt_token + + # Simulate the connect_error handler being called with "User not found" + # This simulates what happens when the server restarts and Redis is cleared + client.socket._on_connect_error("User not found. Please register first.") + + # Verify that login was called (re-registration happened) + # The user should still be the same (or similar guest name) + assert client.user is not None + + # The JWT token should have been refreshed + # (In practice it might be the same token if user existed, but the call was made) + assert client.api.jwt_token is not None + + client.disconnect() + + +def test_connect_error_handler_ignores_other_errors(server): + """Test that _on_connect_error doesn't re-register for non-user-not-found errors.""" + room = uuid.uuid4().hex + client = ZnDraw(room=room, url=server) + assert client.socket.connected + + original_user = client.user + + # Mock the login method to track if it gets called + with patch.object(client.api, "login") as mock_login: + # Simulate a different error (not "User not found") + client.socket._on_connect_error("Connection timeout") + + # Login should NOT be called for other errors + mock_login.assert_not_called() + + # User should remain unchanged + assert client.user == original_user + + client.disconnect() + + +def test_connect_error_handler_handles_login_failure_gracefully(server): + """Test that _on_connect_error handles login failure gracefully.""" + room = uuid.uuid4().hex + client = ZnDraw(room=room, url=server) + assert client.socket.connected + + original_user = client.user + + # Mock login to raise an exception + with patch.object( + client.api, "login", side_effect=Exception("Network error") + ) as mock_login: + # Should not raise, just log the error + client.socket._on_connect_error("User not found") + + # Login was attempted + mock_login.assert_called_once() + + # User should remain unchanged (login failed) + assert client.user == original_user + + client.disconnect() + + +def test_socket_reconnect_after_user_reregistration(server): + """Test full reconnection flow after user re-registration. + + This test simulates what happens when: + 1. Client is connected + 2. Server restarts (simulated by clearing user data) + 3. Socket reconnects and gets "User not found" + 4. Client re-registers and reconnects successfully + """ + import redis + + room = uuid.uuid4().hex + client = ZnDraw(room=room, url=server) + assert client.socket.connected + + original_user = client.user + original_session_id = client.api.session_id + + # Get Redis client to manipulate user data + r = redis.Redis(host="localhost", port=6379, decode_responses=True) + + # Simulate server restart by deleting the user data from Redis + user_key = f"users:data:{original_user}" + user_index_key = "users:index" + + # Delete user data (simulates Redis cleared on restart) + r.delete(user_key) + r.srem(user_index_key, original_user) + + # Now trigger the connect_error handler as if reconnection failed + client.socket._on_connect_error("User not found. Please register first.") + + # Verify user was re-registered (user data should exist again) + assert r.exists(user_key) or client.user != original_user + + # User should be set (either same or new guest name) + assert client.user is not None + assert client.api.jwt_token is not None + + client.disconnect() + + +def test_connect_error_with_password_user(server): + """Test that _on_connect_error works correctly with password-authenticated users.""" + from pydantic import SecretStr + + room = uuid.uuid4().hex + + # Create client (will be guest user, no password) + client = ZnDraw(room=room, url=server) + assert client.socket.connected + + # Simulate having a password set (even though guest users don't need it) + client.password = SecretStr("test-password") + + # Mock login to verify password is passed correctly + with patch.object(client.api, "login") as mock_login: + mock_login.return_value = { + "userName": client.user, + "role": "guest", + "token": "mock-token", + } + + client.socket._on_connect_error("User not found") + + # Verify login was called with correct password + mock_login.assert_called_once_with( + user_name=client.user, + password="test-password", + ) + + client.disconnect() + + +def test_connect_error_updates_role_from_login_response(server): + """Test that _on_connect_error updates role from login response.""" + room = uuid.uuid4().hex + client = ZnDraw(room=room, url=server) + assert client.socket.connected + + original_role = client.role + + # Mock login to return a different role + with patch.object(client.api, "login") as mock_login: + mock_login.return_value = { + "userName": "new-admin-user", + "role": "admin", + "token": "mock-token", + } + + client.socket._on_connect_error("User not found") + + # Role should be updated from login response + assert client.role == "admin" + assert client.user == "new-admin-user" + + client.disconnect()