from __future__ import annotations import asyncio import os import sys import threading import unittest from pathlib import Path from unittest.mock import AsyncMock, Mock, patch from websockets.legacy.server import serve BACKEND_DIR = Path(__file__).resolve().parents[1] if str(BACKEND_DIR) not in sys.path: sys.path.insert(0, str(BACKEND_DIR)) from rpa_engine.douyin_im import ws_client as ws_module from rpa_engine.douyin_im.session import DouyinImSession from rpa_engine.douyin_im.ws_client import DouyinImWsClient, _reconnect_delay TEST_WS_URL = "wss://frontier-im.douyin.com/ws/v2?token=test-token-value" class _FakeWebSocket: def __init__(self, frames=()): self.frames = list(frames) self.next_calls = 0 self.close_code = 1000 self.close_reason = "test complete" self.close_calls: list[tuple[int, str]] = [] self.fail_calls = 0 async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, traceback): return False def __aiter__(self): return self async def __anext__(self): self.next_calls += 1 if not self.frames: raise StopAsyncIteration return self.frames.pop(0) async def close(self, code=1000, reason=""): self.close_calls.append((code, reason)) def fail_connection(self): self.fail_calls += 1 class WebSocketScalingTests(unittest.IsolatedAsyncioTestCase): def _make_client(self, handler=None, account_id=23): session = DouyinImSession( cookies={"sessionid": "session-value", "sid_tt": "sid-value"}, ws_urls=[TEST_WS_URL], user_agent="test-agent/1.0", ) return DouyinImWsClient( session, handler or AsyncMock(), account_id=account_id, ) async def test_async_connection_preserves_handshake_and_ping_options(self): received: list[bytes] = [] async def handler(item): received.append(item["payload"]) client = self._make_client(handler) client._running = True fake_ws = _FakeWebSocket([b"binary-frame", "text-frame"]) connect_mock = Mock(return_value=fake_ws) with ( patch.object(ws_module, "websocket_connect", connect_mock), patch.object( ws_module, "parse_ws_payload", side_effect=lambda payload: [{"payload": payload}], ), patch.object(ws_module.system_logger, "record"), ): await client._run_connection(TEST_WS_URL) queue = client._message_queue self.assertIsNotNone(queue) await asyncio.wait_for(queue.join(), timeout=0.5) await client.stop() self.assertEqual(received, [b"binary-frame", b"text-frame"]) kwargs = connect_mock.call_args.kwargs self.assertEqual(kwargs["origin"], "https://www.douyin.com") self.assertEqual(kwargs["subprotocols"], ["binary", "base64", "pbbp2"]) self.assertEqual(kwargs["user_agent_header"], "test-agent/1.0") self.assertEqual(kwargs["ping_interval"], 20) self.assertEqual(kwargs["ping_timeout"], ws_module._PING_TIMEOUT_SECONDS) self.assertEqual(kwargs["max_queue"], ws_module._TRANSPORT_MAX_QUEUE) self.assertEqual(kwargs["max_size"], ws_module._INCOMING_MAX_SIZE) headers = dict(kwargs["extra_headers"]) self.assertEqual(headers["Cookie"], "sessionid=session-value; sid_tt=sid-value") self.assertNotIn("Sec-WebSocket-Protocol", headers) self.assertFalse(client.connected) self.assertIsNone(client._connection) async def test_starting_500_clients_does_not_create_os_threads(self): parked = asyncio.Event() async def parked_run_loop(_client, _url): await parked.wait() clients = [self._make_client(account_id=index + 1) for index in range(500)] before_threads = threading.active_count() with patch.object(DouyinImWsClient, "_run_loop", parked_run_loop): await asyncio.gather(*(client.start() for client in clients)) await asyncio.sleep(0) self.assertEqual(threading.active_count(), before_threads) self.assertEqual(sum(client._task is not None for client in clients), 500) self.assertEqual( sum(client._dispatcher_task is not None for client in clients), 500, ) await asyncio.gather(*(client.stop() for client in clients)) self.assertTrue(all(client._task is None for client in clients)) self.assertTrue(all(client._dispatcher_task is None for client in clients)) async def test_global_handler_concurrency_is_shared_across_clients(self): limit = 3 release_handlers = asyncio.Event() limit_reached = asyncio.Event() active = 0 maximum_active = 0 started = 0 completed = 0 async def handler(_item): nonlocal active, maximum_active, started, completed active += 1 started += 1 maximum_active = max(maximum_active, active) if started == limit: limit_reached.set() try: await release_handlers.wait() finally: active -= 1 completed += 1 clients = [ self._make_client(handler, account_id=index + 1000) for index in range(12) ] queues = [] with patch.dict( os.environ, {ws_module._HANDLER_CONCURRENCY_ENV: str(limit)}, ): try: for index, client in enumerate(clients): client._running = True client._ensure_dispatcher() queue = client._message_queue self.assertIsNotNone(queue) queues.append(queue) await queue.put({"index": index}) await asyncio.wait_for(limit_reached.wait(), timeout=0.5) # Give every other account a chance to contend for the same # process/event-loop-wide semaphore. await asyncio.sleep(0) self.assertEqual(started, limit) self.assertEqual(maximum_active, limit) release_handlers.set() await asyncio.wait_for( asyncio.gather(*(queue.join() for queue in queues)), timeout=1.0, ) finally: release_handlers.set() await asyncio.gather(*(client.stop() for client in clients)) self.assertEqual(completed, len(clients)) self.assertEqual(maximum_active, limit) async def test_global_handler_limit_preserves_single_client_fifo(self): received: list[int] = [] async def handler(item): await asyncio.sleep(0) received.append(item["sequence"]) client = self._make_client(handler, account_id=2001) client._running = True client._ensure_dispatcher() queue = client._message_queue self.assertIsNotNone(queue) for sequence in range(20): await queue.put({"sequence": sequence}) await asyncio.wait_for(queue.join(), timeout=0.5) await client.stop() self.assertEqual(received, list(range(20))) async def test_stop_cancels_dispatcher_waiting_for_global_handler_slot(self): holder_started = asyncio.Event() release_holder = asyncio.Event() waiter_handler = AsyncMock() async def holder_handler(_item): holder_started.set() await release_holder.wait() holder = self._make_client(holder_handler, account_id=3001) waiter = self._make_client(waiter_handler, account_id=3002) with patch.dict( os.environ, {ws_module._HANDLER_CONCURRENCY_ENV: "1"}, ): holder._running = True holder._ensure_dispatcher() holder_queue = holder._message_queue self.assertIsNotNone(holder_queue) await holder_queue.put({"id": "holder"}) await asyncio.wait_for(holder_started.wait(), timeout=0.5) waiter._running = True waiter._ensure_dispatcher() waiter_queue = waiter._message_queue self.assertIsNotNone(waiter_queue) await waiter_queue.put({"id": "waiter"}) state = ws_module._get_loop_state() async def wait_until_slot_has_waiter(): while not state.handler_slots._waiters: await asyncio.sleep(0) await asyncio.wait_for(wait_until_slot_has_waiter(), timeout=0.5) joined = asyncio.create_task(waiter_queue.join()) await asyncio.wait_for(waiter.stop(), timeout=0.5) await asyncio.wait_for(joined, timeout=0.5) waiter_handler.assert_not_awaited() self.assertTrue(waiter_queue.empty()) release_holder.set() await asyncio.wait_for(holder_queue.join(), timeout=0.5) await holder.stop() async def test_bounded_dispatch_queue_keeps_receiver_responsive(self): first_handler_started = asyncio.Event() release_first_handler = asyncio.Event() received: list[bytes] = [] async def handler(item): received.append(item["payload"]) if len(received) == 1: first_handler_started.set() await release_first_handler.wait() client = self._make_client(handler) client._running = True fake_ws = _FakeWebSocket([b"first", b"second", b"third", b"fourth"]) async def wait_until_third_frame_is_read(): while fake_ws.next_calls < 3: await asyncio.sleep(0) with ( patch.object(ws_module, "_APPLICATION_QUEUE_SIZE", 1), patch.object(ws_module, "websocket_connect", return_value=fake_ws), patch.object( ws_module, "parse_ws_payload", side_effect=lambda payload: [{"payload": payload}], ), patch.object(ws_module.system_logger, "record"), ): task = asyncio.create_task(client._run_connection(TEST_WS_URL)) await asyncio.wait_for(first_handler_started.wait(), timeout=0.5) await asyncio.wait_for(wait_until_third_frame_is_read(), timeout=0.5) # The receiver keeps consuming while the business handler is # blocked, but stops after the bounded application queue fills. # It has pulled the third frame and is blocked enqueueing it; the # fourth frame hasn't been requested and memory remains bounded. self.assertEqual(fake_ws.next_calls, 3) self.assertEqual(received, [b"first"]) self.assertEqual(client._message_queue.qsize(), 1) self.assertIsNotNone(client._dispatcher_task) release_first_handler.set() await asyncio.wait_for(task, timeout=0.5) queue = client._message_queue self.assertIsNotNone(queue) await asyncio.wait_for(queue.join(), timeout=0.5) await client.stop() self.assertEqual(received, [b"first", b"second", b"third", b"fourth"]) self.assertEqual(fake_ws.next_calls, 5) # four frames + end-of-stream async def test_real_async_handshake_receives_binary_frame(self): received: list[bytes] = [] request: dict[str, str | None] = {} async def handler(item): received.append(item["payload"]) async def server_handler(websocket, _path): request["origin"] = websocket.request_headers.get("Origin") request["cookie"] = websocket.request_headers.get("Cookie") request["user_agent"] = websocket.request_headers.get("User-Agent") request["subprotocol"] = websocket.subprotocol await websocket.send(b"protobuf-frame") await websocket.close(code=1000, reason="test complete") client = self._make_client(handler) client._running = True with ( patch.object( ws_module, "parse_ws_payload", side_effect=lambda payload: [{"payload": payload}], ), patch.object(ws_module.system_logger, "record"), ): async with serve( server_handler, "127.0.0.1", 0, origins=["https://www.douyin.com"], subprotocols=["pbbp2"], ) as server: port = server.sockets[0].getsockname()[1] await asyncio.wait_for( client._run_connection(f"ws://127.0.0.1:{port}"), timeout=1.0, ) queue = client._message_queue self.assertIsNotNone(queue) await asyncio.wait_for(queue.join(), timeout=0.5) await client.stop() self.assertEqual(received, [b"protobuf-frame"]) self.assertEqual(request["origin"], "https://www.douyin.com") self.assertEqual(request["cookie"], "sessionid=session-value; sid_tt=sid-value") self.assertEqual(request["user_agent"], "test-agent/1.0") self.assertEqual(request["subprotocol"], "pbbp2") async def test_stop_cancels_slow_handler_and_drains_pending_messages(self): handler_started = asyncio.Event() handler_cancelled = asyncio.Event() never_release = asyncio.Event() async def handler(_item): handler_started.set() try: await never_release.wait() except asyncio.CancelledError: handler_cancelled.set() raise client = self._make_client(handler) client._running = True client._ensure_dispatcher() queue = client._message_queue self.assertIsNotNone(queue) await queue.put({"id": 1}) await queue.put({"id": 2}) await asyncio.wait_for(handler_started.wait(), timeout=0.5) joined = asyncio.create_task(queue.join()) await asyncio.wait_for(client.stop(), timeout=0.5) await asyncio.wait_for(joined, timeout=0.5) self.assertTrue(handler_cancelled.is_set()) self.assertTrue(queue.empty()) self.assertIsNone(client._message_queue) self.assertIsNone(client._dispatcher_task) async def test_first_connection_uses_captured_url_without_refresh(self): client = self._make_client(account_id=24) client._running = True refreshed_url = TEST_WS_URL + "&refreshed=1" client._prepare_url = AsyncMock(return_value=refreshed_url) connected_urls: list[str] = [] async def connection(url): connected_urls.append(url) client._last_connection_lifetime = 1.0 if len(connected_urls) == 2: client._running = False client._run_connection = connection with ( patch.object(ws_module, "_reconnect_delay", return_value=0.0), patch.object(ws_module.system_logger, "record"), ): await asyncio.wait_for(client._run_loop(TEST_WS_URL), timeout=0.5) self.assertEqual(connected_urls, [TEST_WS_URL, refreshed_url]) client._prepare_url.assert_awaited_once_with(TEST_WS_URL) async def test_stop_closes_connection_and_cancels_receive_task(self): client = self._make_client() client._running = True client.connected = True fake_ws = _FakeWebSocket() client._connection = fake_ws task = asyncio.create_task(asyncio.sleep(30)) client._task = task await client.stop() self.assertEqual(fake_ws.close_calls, [(1000, "client stopping")]) self.assertTrue(task.cancelled()) self.assertFalse(client.connected) self.assertIsNone(client._connection) self.assertIsNone(client._task) async def test_stop_aborts_connection_when_close_handshake_stalls(self): client = self._make_client() client._running = True client.connected = True fake_ws = _FakeWebSocket() async def stalled_close(code=1000, reason=""): fake_ws.close_calls.append((code, reason)) await asyncio.Event().wait() fake_ws.close = stalled_close client._connection = fake_ws task = asyncio.create_task(asyncio.sleep(30)) client._task = task with patch.object(ws_module, "_CLOSE_GRACE_SECONDS", 0.01): await asyncio.wait_for(client.stop(), timeout=0.2) self.assertEqual(fake_ws.fail_calls, 1) self.assertTrue(task.cancelled()) self.assertIsNone(client._connection) async def test_short_normal_closes_continue_exponential_retry(self): client = self._make_client(account_id=41) client._running = True client._prepare_url = AsyncMock(return_value=TEST_WS_URL) attempts = 0 async def short_connection(_url): nonlocal attempts attempts += 1 client._last_connection_lifetime = 1.0 if attempts == 4: client._running = False client._run_connection = short_connection with ( patch.object(ws_module, "_reconnect_delay", return_value=0.0) as delay, patch.object(ws_module.system_logger, "record"), ): await asyncio.wait_for(client._run_loop(TEST_WS_URL), timeout=0.5) self.assertEqual([call.args[1] for call in delay.call_args_list], [1, 2, 3]) async def test_stable_connection_resets_retry_counter(self): client = self._make_client(account_id=42) client._running = True client._prepare_url = AsyncMock(return_value=TEST_WS_URL) lifetimes = [1.0, 1.0, ws_module._STABLE_CONNECTION_SECONDS + 1.0, 1.0] attempts = 0 async def connection(_url): nonlocal attempts client._last_connection_lifetime = lifetimes[attempts] attempts += 1 if attempts == len(lifetimes): client._running = False client._run_connection = connection with ( patch.object(ws_module, "_reconnect_delay", return_value=0.0) as delay, patch.object(ws_module.system_logger, "record"), ): await asyncio.wait_for(client._run_loop(TEST_WS_URL), timeout=0.5) self.assertEqual([call.args[1] for call in delay.call_args_list], [1, 2, 1]) async def test_repeated_connection_failures_throttle_system_logs_per_account(self): client = self._make_client(account_id=4041) client._running = True client._prepare_url = AsyncMock(return_value=TEST_WS_URL) attempts = 0 async def failing_connection(_url): nonlocal attempts attempts += 1 if attempts == 6: client._running = False raise ConnectionError("frontier unavailable") client._run_connection = failing_connection with ( patch.object(ws_module, "_reconnect_delay", return_value=0.0), patch.object( ws_module, "_system_log_throttle_seconds", return_value=300.0, ), patch.object(ws_module.system_logger, "record") as system_record, patch.object(ws_module.logger, "warning") as ordinary_warning, ): await asyncio.wait_for(client._run_loop(TEST_WS_URL), timeout=0.5) # Ordinary diagnostics remain available for every live failure, while # the database-backed system log receives one row for the repeated # failure/reconnect cycle of this account. self.assertEqual(attempts, 6) self.assertEqual(ordinary_warning.call_count, 5) self.assertEqual(system_record.call_count, 1) self.assertEqual(system_record.call_args.kwargs["account_id"], 4041) def test_handler_concurrency_environment_value_is_safely_clamped(self): with patch.dict(os.environ, {ws_module._HANDLER_CONCURRENCY_ENV: "0"}): self.assertEqual(ws_module._handler_concurrency_limit(), 1) with patch.dict(os.environ, {ws_module._HANDLER_CONCURRENCY_ENV: "500"}): self.assertEqual(ws_module._handler_concurrency_limit(), 32) with patch.dict(os.environ, {ws_module._HANDLER_CONCURRENCY_ENV: "invalid"}): self.assertEqual(ws_module._handler_concurrency_limit(), 8) def test_reconnect_delay_is_bounded_and_spread_by_account(self): waits = [_reconnect_delay(10, retry) for retry in range(1, 10)] self.assertGreater(waits[1], waits[0]) self.assertGreater(waits[2], waits[1]) self.assertTrue(all(2.0 <= wait < 90.0 for wait in waits)) self.assertNotEqual(_reconnect_delay(10, 8), _reconnect_delay(11, 8)) if __name__ == "__main__": unittest.main()