import sys import os import asyncio import json import unittest from unittest.mock import MagicMock, AsyncMock, patch # Add project root to path sys.path.append(os.getcwd()) # Mock app.config BEFORE imports to prevent attempt to load real settings mock_config = MagicMock() mock_config.settings = MagicMock() mock_config.settings.REDIS = {'server': 'localhost', 'port': 6379} sys.modules["app.config"] = mock_config # Mock DB and circular/heavy imports before any app module is loaded sys.modules["app.db_sql"] = MagicMock() sys.modules["app.lib_sql_core"] = MagicMock() sys.modules["app.db_connection"] = MagicMock() sys.modules["app.log"] = MagicMock() # Provide a real AccountContext but stub out the heavy lib_general_v3 imports from app.models.auth_models import AccountContext mock_lib_general_v3 = MagicMock() mock_lib_general_v3.AccountContext = AccountContext mock_lib_general_v3.get_account_context_optional = MagicMock(return_value=AccountContext( account_id=1, account_id_random="test_account_id", auth_method="api_key" )) sys.modules["app.lib_general_v3"] = mock_lib_general_v3 sys.modules["app.routers.dependencies_v3"] = MagicMock() sys.modules["app.lib_jwt"] = MagicMock() from app.routers.websockets_v3 import v3_ws_endpoint # Shared AccountContext fixture used in all test calls MOCK_ACCOUNT = AccountContext(account_id=1, account_id_random="test_account_id", auth_method="api_key") class TestWSV3Router(unittest.TestCase): @patch('app.routers.websockets_v3.ws_manager_v3') def test_v3_ws_endpoint_logic(self, mock_manager): """ Tests the core logic of the V3 WebSocket endpoint, ensuring Redis subscription and bidirectional message handling are initiated. """ # 1. Setup WebSocket Mock mock_ws = AsyncMock() # 2. Setup Redis PubSub Mock mock_pubsub = MagicMock() mock_pubsub.subscribe = AsyncMock() mock_pubsub.unsubscribe = AsyncMock() mock_pubsub.close = AsyncMock() mock_message = { 'type': 'message', 'data': json.dumps({ "version": "3", "msg_type": "msg", "target": "group", "from_id": "other_client", "msg": "Hello from Redis", "payload": {}, "sent_at": "2026-01-30T12:00:00Z" }) } # Signal to coordinate loops msg_delivered = asyncio.Event() # Counters to break the 'while True' loops in the endpoint get_msg_count = 0 recv_json_count = 0 async def mock_get_message(*args, **kwargs): nonlocal get_msg_count get_msg_count += 1 if get_msg_count == 1: msg_delivered.set() return mock_message await asyncio.sleep(0.05) # Raise CancelledError to terminate the loop cleanly raise asyncio.CancelledError("Terminate sender loop") mock_pubsub.get_message = mock_get_message mock_redis = MagicMock() mock_redis.pubsub.return_value = mock_pubsub # 3. Setup Manager Mock mock_manager.get_redis = AsyncMock(return_value=mock_redis) mock_manager.update_presence = AsyncMock() mock_manager.publish_message = AsyncMock() mock_manager.get_channel_names.return_value = ["ws:group:test"] # Mock incoming websocket message async def mock_receive_json(): nonlocal recv_json_count recv_json_count += 1 if recv_json_count == 1: # Wait until the sender loop has processed the Redis message await msg_delivered.wait() return { "msg_type": "msg", "target": "group", "msg": "Client A saying hi" } await asyncio.sleep(0.05) # Raise CancelledError to terminate the loop cleanly raise asyncio.CancelledError("Terminate receiver loop") mock_ws.receive_json.side_effect = mock_receive_json # 4. Run the endpoint logic async def run_endpoint(): try: await asyncio.wait_for( v3_ws_endpoint(mock_ws, "test_group", "client_a", account=MOCK_ACCOUNT), timeout=0.5 ) except (asyncio.TimeoutError, asyncio.CancelledError): pass except Exception as e: if "Terminate" not in str(e): raise asyncio.run(run_endpoint()) # 5. Verifications mock_ws.accept.assert_called_once() mock_manager.update_presence.assert_any_call("client_a", "test_group", online=True) mock_ws.send_text.assert_called() print("✅ WebSocket Router unit logic verified.") @patch('app.routers.websockets_v3.ws_manager_v3') def test_heartbeat_refreshes_presence(self, mock_manager): """ Ensures that a heartbeat message triggers a presence TTL refresh in addition to being published normally. """ mock_ws = AsyncMock() mock_pubsub = MagicMock() mock_pubsub.subscribe = AsyncMock() mock_pubsub.unsubscribe = AsyncMock() mock_pubsub.close = AsyncMock() async def mock_get_message(*args, **kwargs): await asyncio.sleep(0.05) raise asyncio.CancelledError("Terminate sender loop") mock_pubsub.get_message = mock_get_message mock_redis = MagicMock() mock_redis.pubsub.return_value = mock_pubsub recv_count = 0 async def mock_receive_json(): nonlocal recv_count recv_count += 1 if recv_count == 1: # Send a heartbeat return {"msg_type": "heartbeat", "target": "echo"} await asyncio.sleep(0.05) raise asyncio.CancelledError("Terminate receiver loop") mock_ws.receive_json.side_effect = mock_receive_json mock_manager.get_redis = AsyncMock(return_value=mock_redis) mock_manager.update_presence = AsyncMock() mock_manager.publish_message = AsyncMock() mock_manager.get_channel_names.return_value = ["ws:group:test"] async def run_endpoint(): try: await asyncio.wait_for( v3_ws_endpoint(mock_ws, "test_group", "client_hb", account=MOCK_ACCOUNT), timeout=0.5 ) except (asyncio.TimeoutError, asyncio.CancelledError): pass except Exception as e: if "Terminate" not in str(e): raise asyncio.run(run_endpoint()) # Presence should be refreshed: once on connect, once on heartbeat receipt presence_calls = mock_manager.update_presence.call_args_list online_calls = [c for c in presence_calls if c.kwargs.get('online') is True or (len(c.args) >= 3 and c.args[2] is True)] self.assertGreaterEqual(len(online_calls), 2, "Expected at least 2 presence refresh calls (connect + heartbeat)") print("✅ Heartbeat presence refresh verified.")