dy/backend/tests/test_ws_client_scaling.py

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()