574 lines
22 KiB
Python
574 lines
22 KiB
Python
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?fpid=9&device_id=10001&"
|
|
"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_browser_frontier_uses_text_heartbeat_and_filters_ack(self):
|
|
client = self._make_client()
|
|
client._running = True
|
|
websocket = AsyncMock()
|
|
|
|
heartbeat = asyncio.create_task(client._run_browser_heartbeat(websocket))
|
|
for _ in range(20):
|
|
if websocket.send.await_count:
|
|
break
|
|
await asyncio.sleep(0)
|
|
heartbeat.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await heartbeat
|
|
|
|
websocket.send.assert_awaited_once_with("hi")
|
|
with patch.object(ws_module, "parse_ws_payload") as parse:
|
|
await client._dispatch("hi")
|
|
parse.assert_not_called()
|
|
|
|
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()
|