""" Integration tests for WebSocket handlers. This module tests the full WebSocket connection lifecycle, message handling, rate limiting, and error handling according to TESTING_PLAN.md. Tests: Connection Lifecycle: - Connection acceptance and initialization - Connection rejection when max connections reached - Disconnection cleanup - Reconnection and game loading - Multiple connections from same user Message Handling: - Init message processing - Command message processing (start/stop/restart) - Speed change messages - Question responses - Conversation messages - Invalid message types - Malformed JSON Rate Limiting: - Rate limit enforcement - Rate limit reset after time window - Per-user rate limiting Error Handling: - Error isolation (doesn't crash server) - Connection timeout - Ping/pong keepalive """ import sys import os sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))) import pytest import asyncio import json from unittest.mock import Mock, AsyncMock, patch, MagicMock from datetime import datetime, timedelta # Import WebSocket handlers and related modules from server.websocket_handlers import start, shutdown, error, handler from server.websocket_registry import UserRegistry, USERS from player_cache import PlayerCache from functions import playerClass, personClass, getOccupations from config import config from rate_limiter import RateLimiter # ============================================================================ # Fixtures # ============================================================================ @pytest.fixture def mock_websocket(): """Create a mock WebSocket connection.""" ws = AsyncMock() ws.userID = 'test_user_123' ws.send = AsyncMock() ws.recv = AsyncMock() ws.close = AsyncMock() ws.closed = False return ws @pytest.fixture def mock_websocket_different_user(): """Create a second mock WebSocket for multi-user testing.""" ws = AsyncMock() ws.userID = 'test_user_456' ws.send = AsyncMock() ws.recv = AsyncMock() ws.close = AsyncMock() ws.closed = False return ws @pytest.fixture def mock_player(): """Create a mock player object.""" player = playerClass() player.userID = 'test_user_123' player.id = 'test_user_123' player.c = personClass() player.c.firstname = 'Test' player.c.lastname = 'User' player.c.sex = 'Male' player.c.ageYears = 25 player.c.energy = 100 player.c.money = 1000 player.c.happiness = 75 player.date = datetime.now() player.hourOfDay = 12 player.minuteOfHour = 0 player.gameSpeed = config.SPEED_DEFAULT player.controller = 'inactive' player.connection = 'disconnected' player.occupations = getOccupations() return player @pytest.fixture def player_cache(): """Create a fresh player cache for each test.""" return PlayerCache(max_size=10) @pytest.fixture def user_registry(): """Create a fresh user registry for each test.""" return UserRegistry() @pytest.fixture(autouse=True) def setup_app_module(player_cache, monkeypatch): """Mock the app module's playerRecords.""" import app monkeypatch.setattr(app, 'playerRecords', player_cache) yield # Cleanup player_cache._cache.clear() @pytest.fixture def mock_db_operations(): """Mock database operations.""" with patch('functions.loadGameAsync', new_callable=AsyncMock) as mock_load, \ patch('functions.saveGameAsync', new_callable=AsyncMock) as mock_save, \ patch('functions.insertGame', new_callable=Mock) as mock_insert: mock_load.return_value = None # No saved game by default mock_save.return_value = True mock_insert.return_value = True yield { 'load': mock_load, 'save': mock_save, 'insert': mock_insert } # ============================================================================ # Connection Lifecycle Tests # ============================================================================ class TestConnectionLifecycle: """Tests for WebSocket connection lifecycle management.""" @pytest.mark.asyncio async def test_websocket_connection_accepted(self, mock_websocket, mock_db_operations, user_registry): """Test that a WebSocket connection is accepted and initialized.""" # Setup init_message = json.dumps({"type": "init", "userID": "test_user_123"}) mock_websocket.recv.return_value = init_message # Mock the start function to avoid full initialization with patch('server.websocket_handlers.start', new_callable=AsyncMock) as mock_start: mock_start.return_value = None # Execute try: await handler(mock_websocket) except: pass # Handler closes connection at the end # Verify mock_websocket.recv.assert_called_once() mock_start.assert_called_once_with(mock_websocket) assert mock_websocket.userID == 'test_user_123' @pytest.mark.asyncio async def test_websocket_connection_rejected_if_max_reached(self, mock_websocket, player_cache): """Test that cache doesn't evict connected players when full.""" # Setup: Fill the cache to max capacity with connected players player_cache._max_size = 2 for i in range(2): test_player = playerClass() test_player.userID = f'user_{i}' test_player.id = f'user_{i}' test_player.c = personClass() test_player.c.firstname = f'User{i}' test_player.connection = 'connected' player_cache.set(f'user_{i}', test_player) # Verify cache is full assert player_cache.size() == 2 # Try to add another connection when all existing are connected new_player = playerClass() new_player.userID = 'user_overflow' new_player.id = 'user_overflow' new_player.c = personClass() new_player.c.firstname = 'Overflow' new_player.connection = 'connected' # Cache should NOT evict connected players (logs warning instead) # So size will exceed max_size temporarily player_cache.set('user_overflow', new_player) # All 3 players should still be in cache (no eviction of connected players) assert player_cache.size() == 3 @pytest.mark.asyncio async def test_websocket_disconnection_cleanup(self, mock_websocket, mock_player, player_cache, user_registry, mock_db_operations): """Test that disconnection properly cleans up resources.""" # Setup mock_websocket.userID = mock_player.userID player_cache.set(mock_player.userID, mock_player) user_registry.add(mock_websocket) mock_player.connection = 'connected' mock_player.controller = 'active' # Patch saveGameAsync and USERS where they're used in websocket_handlers with patch('server.websocket_handlers.saveGameAsync', new_callable=AsyncMock) as mock_save, \ patch('server.websocket_handlers.USERS', user_registry): mock_save.return_value = True # Execute result = await shutdown(mock_websocket) # Verify assert result is True assert mock_player.connection == 'disconnected' assert mock_player.controller == 'inactive' assert mock_player.offlineStats.minutesOffline == 0 mock_save.assert_called_once_with(mock_player) assert user_registry.get(mock_player.userID) is None @pytest.mark.asyncio async def test_websocket_reconnection_loads_game(self, mock_websocket, mock_player, player_cache, mock_db_operations): """Test that reconnection loads the player's saved game.""" # Setup: Mock loadGameAsync to return a saved player mock_player.gameSpeed = config.SPEED_QUESTION_PAUSE # Paused state with patch('server.websocket_handlers.loadGameAsync', new_callable=AsyncMock) as mock_load, \ patch('server.websocket_handlers.sendUserInfo', new_callable=AsyncMock) as mock_send, \ patch('retention.daily_rewards.handle_daily_login_check') as mock_daily, \ patch('server.websocket_handlers.insertGame', new_callable=Mock) as mock_insert, \ patch('server.websocket_handlers.connect') as mock_connect, \ patch('asyncio.gather', new_callable=AsyncMock): mock_load.return_value = mock_player # Execute await start(mock_websocket) # Verify mock_load.assert_called_once_with('test_user_123') assert mock_player.controller == 'active' assert mock_player.connection == 'connected' assert mock_player.gameSpeed == config.SPEED_DEFAULT # Should reset from paused mock_send.assert_called_once() @pytest.mark.asyncio async def test_websocket_multiple_connections_same_user(self, mock_websocket, mock_websocket_different_user, player_cache, user_registry): """Test handling multiple connections from the same user (different devices).""" # Setup same_user_id = 'test_user_999' mock_websocket.userID = same_user_id mock_websocket_different_user.userID = same_user_id # First connection user_registry.add(mock_websocket) assert user_registry.get(same_user_id) == mock_websocket # Second connection (should replace first - upsert behavior) user_registry.add(mock_websocket_different_user) assert user_registry.get(same_user_id) == mock_websocket_different_user # Only one connection should be tracked assert user_registry.count() == 1 # ============================================================================ # Message Handling Tests # ============================================================================ class TestMessageHandling: """Tests for WebSocket message processing.""" @pytest.mark.asyncio async def test_websocket_receives_init_message(self, mock_websocket, mock_db_operations): """Test processing of init message.""" # Setup init_message = json.dumps({"type": "init", "userID": "test_user_123"}) mock_websocket.recv.return_value = init_message with patch('server.websocket_handlers.start', new_callable=AsyncMock) as mock_start: mock_start.return_value = None try: await handler(mock_websocket) except: pass # Verify init message was processed assert mock_websocket.userID == 'test_user_123' mock_start.assert_called_once() @pytest.mark.asyncio async def test_websocket_receives_command_message(self, mock_websocket, mock_player, player_cache): """Test processing of command messages (start/stop).""" from game_loop.producer_consumer import consumer # Setup player_cache.set(mock_player.userID, mock_player) mock_websocket.userID = mock_player.userID # Test start command start_message = json.dumps({"type": "command", "value": "start"}) with patch('server.command_dispatcher.dispatch_command', new_callable=AsyncMock) as mock_dispatch, \ patch('server.websocket_messaging.sendDict', new_callable=AsyncMock): await consumer(start_message, mock_websocket) # Verify command was dispatched mock_dispatch.assert_called_once() call_args = mock_dispatch.call_args[0] assert call_args[0]['type'] == 'command' assert call_args[0]['value'] == 'start' @pytest.mark.asyncio async def test_websocket_receives_speed_message(self, mock_websocket, mock_player, player_cache): """Test processing of speed change messages.""" from game_loop.producer_consumer import consumer # Setup player_cache.set(mock_player.userID, mock_player) mock_websocket.userID = mock_player.userID speed_message = json.dumps({"type": "speed", "value": 500}) with patch('server.command_dispatcher.dispatch_command', new_callable=AsyncMock) as mock_dispatch, \ patch('server.websocket_messaging.sendDict', new_callable=AsyncMock): await consumer(speed_message, mock_websocket) # Verify mock_dispatch.assert_called_once() call_args = mock_dispatch.call_args[0] assert call_args[0]['type'] == 'speed' assert call_args[0]['value'] == 500 @pytest.mark.asyncio async def test_websocket_receives_question_response(self, mock_websocket, mock_player, player_cache): """Test processing of question response messages.""" from game_loop.producer_consumer import consumer # Setup player_cache.set(mock_player.userID, mock_player) mock_websocket.userID = mock_player.userID question_response = json.dumps({ "type": "questionEvent", "id": "question_123", "response": "option_1" }) with patch('server.command_dispatcher.dispatch_command', new_callable=AsyncMock) as mock_dispatch, \ patch('server.websocket_messaging.sendDict', new_callable=AsyncMock): await consumer(question_response, mock_websocket) # Verify mock_dispatch.assert_called_once() call_args = mock_dispatch.call_args[0] assert call_args[0]['type'] == 'questionEvent' assert call_args[0]['id'] == 'question_123' @pytest.mark.asyncio async def test_websocket_receives_conversation_message(self, mock_websocket, mock_player, player_cache): """Test processing of conversation messages.""" from game_loop.producer_consumer import consumer # Setup player_cache.set(mock_player.userID, mock_player) mock_websocket.userID = mock_player.userID conversation_message = json.dumps({ "type": "conversation", "message": { "id": "conversation_123", "message": "Hello there!" } }) with patch('server.command_dispatcher.dispatch_command', new_callable=AsyncMock) as mock_dispatch, \ patch('server.websocket_messaging.sendDict', new_callable=AsyncMock): await consumer(conversation_message, mock_websocket) # Verify mock_dispatch.assert_called_once() call_args = mock_dispatch.call_args[0] assert call_args[0]['type'] == 'conversation' assert 'message' in call_args[0] @pytest.mark.asyncio async def test_websocket_receives_invalid_message_type(self, mock_websocket, mock_player, player_cache): """Test that invalid message types are handled gracefully.""" from game_loop.producer_consumer import consumer # Setup player_cache.set(mock_player.userID, mock_player) mock_websocket.userID = mock_player.userID invalid_message = json.dumps({ "type": "invalid_command_xyz", "data": "should be ignored" }) with patch('server.command_dispatcher.dispatch_command', new_callable=AsyncMock) as mock_dispatch, \ patch('server.websocket_messaging.sendDict', new_callable=AsyncMock): # Should not raise an exception await consumer(invalid_message, mock_websocket) # Command dispatcher should still be called (it handles unknown commands) mock_dispatch.assert_called_once() @pytest.mark.asyncio async def test_websocket_receives_malformed_json(self, mock_websocket, mock_player, player_cache): """Test that malformed JSON is handled gracefully.""" from game_loop.producer_consumer import consumer # Setup player_cache.set(mock_player.userID, mock_player) mock_websocket.userID = mock_player.userID malformed_message = "{invalid json: not closed properly" # Should handle the error gracefully with pytest.raises(json.JSONDecodeError): await consumer(malformed_message, mock_websocket) # ============================================================================ # Rate Limiting Tests # ============================================================================ class TestRateLimiting: """Tests for WebSocket rate limiting.""" @pytest.mark.asyncio async def test_websocket_rate_limiting_enforced(self, mock_websocket, mock_player, player_cache): """Test that rate limiting is enforced after too many messages.""" from game_loop.producer_consumer import consumer_handler # Setup player_cache.set(mock_player.userID, mock_player) mock_websocket.userID = mock_player.userID # Create a rate limiter with very low limits for testing rate_limiter = RateLimiter(max_requests=3, window_seconds=60) # Send 3 requests (should succeed) for i in range(3): assert rate_limiter.is_allowed(mock_player.userID) is True # 4th request should be rate limited assert rate_limiter.is_allowed(mock_player.userID) is False # Further requests should also be blocked assert rate_limiter.is_allowed(mock_player.userID) is False @pytest.mark.asyncio async def test_websocket_rate_limiting_resets(self): """Test that rate limiting resets after time window.""" # Create a rate limiter with 1 second window rate_limiter = RateLimiter(max_requests=2, window_seconds=1) user_id = 'test_user_rate_limit' # Use up the quota assert rate_limiter.is_allowed(user_id) is True assert rate_limiter.is_allowed(user_id) is True assert rate_limiter.is_allowed(user_id) is False # Wait for window to expire await asyncio.sleep(1.1) # Should be allowed again assert rate_limiter.is_allowed(user_id) is True @pytest.mark.asyncio async def test_websocket_rate_limiting_per_user(self): """Test that rate limiting is enforced per user.""" rate_limiter = RateLimiter(max_requests=2, window_seconds=60) user1 = 'user_1' user2 = 'user_2' # User 1 uses quota assert rate_limiter.is_allowed(user1) is True assert rate_limiter.is_allowed(user1) is True assert rate_limiter.is_allowed(user1) is False # Rate limited # User 2 should still have quota (separate limit) assert rate_limiter.is_allowed(user2) is True assert rate_limiter.is_allowed(user2) is True assert rate_limiter.is_allowed(user2) is False # Now user 2 is limited too # ============================================================================ # Error Handling Tests # ============================================================================ class TestErrorHandling: """Tests for WebSocket error handling and resilience.""" @pytest.mark.asyncio async def test_websocket_error_doesnt_crash_server(self, mock_websocket, mock_player, player_cache, mock_db_operations): """Test that an error in one connection doesn't crash the server.""" from game_loop.producer_consumer import consumer # Setup player_cache.set(mock_player.userID, mock_player) mock_websocket.userID = mock_player.userID # Simulate an error during command processing error_message = json.dumps({"type": "command", "value": "start"}) with patch('server.command_dispatcher.dispatch_command', new_callable=AsyncMock) as mock_dispatch: mock_dispatch.side_effect = Exception("Simulated error") # Should handle the error without crashing with pytest.raises(Exception): await consumer(error_message, mock_websocket) @pytest.mark.asyncio async def test_websocket_connection_timeout(self, mock_websocket): """Test that connection timeout is handled properly.""" import websockets.exceptions # Simulate a connection timeout mock_websocket.recv.side_effect = asyncio.TimeoutError("Connection timeout") with patch('server.websocket_handlers.shutdown', new_callable=AsyncMock) as mock_shutdown: try: await handler(mock_websocket) except asyncio.TimeoutError: pass except Exception: pass @pytest.mark.asyncio async def test_websocket_ping_pong(self, mock_websocket): """Test WebSocket ping/pong keepalive mechanism.""" # WebSocket ping/pong is typically handled at the protocol level # This test verifies that the websocket supports ping/pong # Mock ping method mock_websocket.ping = AsyncMock() mock_websocket.pong = AsyncMock() # Simulate ping await mock_websocket.ping() mock_websocket.ping.assert_called_once() # Pong would be automatically sent by the WebSocket library # We just verify the methods exist and can be called # ============================================================================ # Additional Integration Tests # ============================================================================ class TestWebSocketIntegration: """Additional integration tests for complex scenarios.""" @pytest.mark.asyncio async def test_full_connection_lifecycle(self, mock_websocket, mock_db_operations, player_cache, user_registry): """Test full connection lifecycle from connect to disconnect.""" # Setup init_message = json.dumps({"type": "init", "userID": "lifecycle_test"}) mock_websocket.userID = "lifecycle_test" mock_websocket.recv.return_value = init_message with patch('server.websocket_handlers.start', new_callable=AsyncMock) as mock_start, \ patch('server.websocket_handlers.shutdown', new_callable=AsyncMock) as mock_shutdown: # Connect user_registry.add(mock_websocket) assert user_registry.count() == 1 # Disconnect user_registry.remove(mock_websocket) assert user_registry.count() == 0 @pytest.mark.asyncio async def test_message_serialization_deserialization(self, mock_websocket, mock_player, player_cache): """Test that complex message types are properly serialized/deserialized.""" from game_loop.producer_consumer import consumer from server.websocket_messaging import ComplexHandler # Setup player_cache.set(mock_player.userID, mock_player) mock_websocket.userID = mock_player.userID # Create a complex message with various data types complex_message = { "type": "characterSetup", "message": { "name": "Test Character", "age": 25, "sex": "Male", "timestamp": datetime.now().isoformat(), "attributes": { "strength": 10, "intelligence": 15 } } } message_json = json.dumps(complex_message, default=ComplexHandler) with patch('server.command_dispatcher.dispatch_command', new_callable=AsyncMock) as mock_dispatch, \ patch('server.websocket_messaging.sendDict', new_callable=AsyncMock): await consumer(message_json, mock_websocket) # Verify the message was properly deserialized mock_dispatch.assert_called_once() call_args = mock_dispatch.call_args[0] assert call_args[0]['type'] == 'characterSetup' assert 'message' in call_args[0] @pytest.mark.asyncio async def test_concurrent_connections(self, mock_websocket, mock_websocket_different_user, player_cache, user_registry): """Test handling of concurrent connections from different users.""" # Setup user_registry.add(mock_websocket) user_registry.add(mock_websocket_different_user) # Verify both connections are tracked assert user_registry.count() == 2 assert user_registry.get('test_user_123') == mock_websocket assert user_registry.get('test_user_456') == mock_websocket_different_user @pytest.mark.asyncio async def test_error_message_sent_to_client(self, mock_websocket): """Test that error messages are properly sent to client.""" error_text = "Test error message" await error(mock_websocket, error_text) # Verify error was sent mock_websocket.send.assert_called_once() sent_data = json.loads(mock_websocket.send.call_args[0][0]) assert sent_data['type'] == 'error' assert sent_data['message'] == error_text # ============================================================================ # Summary Statistics # ============================================================================ """ Test Summary: ============= Connection Lifecycle Tests: 5 - test_websocket_connection_accepted - test_websocket_connection_rejected_if_max_reached - test_websocket_disconnection_cleanup - test_websocket_reconnection_loads_game - test_websocket_multiple_connections_same_user Message Handling Tests: 7 - test_websocket_receives_init_message - test_websocket_receives_command_message - test_websocket_receives_speed_message - test_websocket_receives_question_response - test_websocket_receives_conversation_message - test_websocket_receives_invalid_message_type - test_websocket_receives_malformed_json Rate Limiting Tests: 3 - test_websocket_rate_limiting_enforced - test_websocket_rate_limiting_resets - test_websocket_rate_limiting_per_user Error Handling Tests: 3 - test_websocket_error_doesnt_crash_server - test_websocket_connection_timeout - test_websocket_ping_pong Additional Integration Tests: 5 - test_full_connection_lifecycle - test_message_serialization_deserialization - test_concurrent_connections - test_error_message_sent_to_client Total Tests: 23 Coverage: - WebSocket connection acceptance and rejection - Game loading and initialization - Disconnection cleanup and resource management - Message routing and command dispatch - Rate limiting enforcement - Error isolation and resilience - Multi-user support - Message serialization/deserialization """