更新
This commit is contained in:
Your Name
@@ -0,0 +1,551 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user