from __future__ import annotations import asyncio import os import sys import unittest from contextlib import asynccontextmanager from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, patch BACKEND_DIR = Path(__file__).resolve().parents[1] os.environ["KEFU_DB_TYPE"] = "sqlite" os.environ["KEFU_DATABASE_URL"] = "" os.environ["KEFU_DB_PATH"] = str(BACKEND_DIR / "kefu.db") if str(BACKEND_DIR) not in sys.path: sys.path.insert(0, str(BACKEND_DIR)) from rpa_engine.douyin_im.http_client import DouyinImHttpClient from rpa_engine.douyin_im.session import DouyinImSession from rpa_engine.douyin_im.service import DouyinImService, _conversation_poll_timing from rpa_engine.douyin_im import service as service_module class ConversationPollBandwidthTests(unittest.IsolatedAsyncioTestCase): def _make_client(self) -> DouyinImHttpClient: return DouyinImHttpClient( DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001), account_id=9, ) async def test_inbox_is_fetched_with_exactly_one_protobuf_request(self): """imapi 只认 protobuf；轮询一次就只该发一个请求。""" client = self._make_client() client.fetch_inbox_messages = AsyncMock(return_value=[]) client._request = AsyncMock() self.assertEqual(await client.get_conversations(), []) client.fetch_inbox_messages.assert_awaited_once() # 不能再退回 JSON 的 /v1/conversation/list：那个请求恒被抖音拒绝。 client._request.assert_not_awaited() async def test_transport_failure_is_recorded_and_returns_empty(self): client = self._make_client() client.fetch_inbox_messages = AsyncMock(side_effect=RuntimeError("boom")) with self.assertLogs("douyin_im.http", level="WARNING"): self.assertEqual(await client.get_conversations(), []) self.assertIn("boom", client.last_error) async def test_inbox_messages_group_into_one_row_per_conversation(self): client = self._make_client() client.fetch_inbox_messages = AsyncMock( return_value=[ { "conversation_id": "0:1:10001:20001", "server_message_id": "700", "conversation_short_id": "555", "message_type": 7, "sender": "20001", "content": '{"text":"旧"}', }, { "conversation_id": "0:1:10001:20001", "server_message_id": "900", "conversation_short_id": "555", "message_type": 7, "sender": "20001", "content": '{"text":"新"}', }, { "conversation_id": "0:1:10001:20002", "server_message_id": "800", "message_type": 7, "sender": "20002", "content": '{"text":"另一个"}', }, ] ) rows = await client.get_conversations(enrich_profiles=False) by_id = {r["conversation_id"]: r for r in rows} self.assertEqual(len(rows), 2) # 同一会话只保留 server_message_id 最大的那条 self.assertEqual(by_id["0:1:10001:20001"]["server_message_id"], "900") self.assertIn("新", by_id["0:1:10001:20001"]["content"]) # peer_uid 由 conversation_id 推导，不能直接取 sender（可能是自己） self.assertEqual(by_id["0:1:10001:20002"]["peer_uid"], "20002") # 顺手缓存 short_id，发送时就不必再 create 一次会话 self.assertEqual( client.session.conv_meta["0:1:10001:20001"]["conversation_short_id"], "555", ) def test_websocket_reconciliation_is_slow_and_http_fallback_stays_fast(self): with patch.dict( os.environ, { "KEFU_WS_RECONCILE_INTERVAL_SECONDS": "120", "KEFU_HTTP_POLL_INTERVAL_SECONDS": "15", }, ): ws_interval, ws_stagger = _conversation_poll_timing(123, True) http_interval, http_stagger = _conversation_poll_timing(123, False) self.assertEqual(ws_interval, 120) self.assertEqual(http_interval, 15) self.assertGreaterEqual(ws_stagger, 0) self.assertLess(ws_stagger, ws_interval) self.assertGreaterEqual(http_stagger, 0) self.assertLess(http_stagger, http_interval) def test_poll_interval_expands_to_the_configured_population_budget(self): with patch.dict( os.environ, { "KEFU_WS_RECONCILE_INTERVAL_SECONDS": "120", "KEFU_HTTP_POLL_INTERVAL_SECONDS": "15", "KEFU_WS_POLL_BUDGET_RPS": "1", "KEFU_HTTP_POLL_BUDGET_RPS": "1", }, ): ws_interval, _ = _conversation_poll_timing( 123, True, population=500, ) http_interval, _ = _conversation_poll_timing( 123, False, population=500, ) self.assertEqual(ws_interval, 500) self.assertEqual(http_interval, 500) def test_initial_unread_concurrency_is_configurable_and_bounded(self): with patch.dict( os.environ, {"KEFU_INITIAL_UNREAD_CONCURRENCY": "4"}, ): self.assertEqual(service_module._initial_unread_concurrency(), 4) with patch.dict( os.environ, {"KEFU_INITIAL_UNREAD_CONCURRENCY": "999"}, ): self.assertEqual(service_module._initial_unread_concurrency(), 8) async def test_service_poll_uses_one_conversation_request_without_unread_probe(self): class _Controller: @asynccontextmanager async def background_slot(self, *_args, **_kwargs): yield class _HttpClient: def __init__(self): self.get_conversations = AsyncMock(return_value=[]) self.conversation_list_unsupported = False async def __aenter__(self): return self async def __aexit__(self, *_args): return False http = _HttpClient() service = DouyinImService( session=DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001), match_reply=AsyncMock(return_value=None), log_fn=AsyncMock(), account_id=9, ) service._index_conversations = AsyncMock() with ( patch.object(service_module, "get_traffic_controller", return_value=_Controller()), patch.object(service_module, "DouyinImHttpClient", return_value=http), ): await service._poll_conversations() http.get_conversations.assert_awaited_once_with(enrich_profiles=False) service._index_conversations.assert_awaited_once_with( [], enrich_profiles=False, ) async def test_poll_only_handles_unread_or_a_genuinely_changed_preview(self): class _Controller: def __init__(self): self.startup_flags = [] @asynccontextmanager async def background_slot(self, *_args, **kwargs): self.startup_flags.append(bool(kwargs.get("startup"))) yield snapshots = [ [ { "conversation_id": "0:1:10001:20001", "peer_uid": "20001", "sender_name": "历史会话", "sender_avatar": "https://example.test/a.png", "content": "历史消息", "unread_count": 0, }, { "conversation_id": "0:1:10001:20002", "peer_uid": "20002", "sender_name": "未读会话", "sender_avatar": "https://example.test/b.png", "content": "新消息", "unread_count": 1, }, ], [ { "conversation_id": "0:1:10001:20001", "peer_uid": "20001", "sender_name": "历史会话", "sender_avatar": "https://example.test/a.png", "content": "历史消息", "unread_count": 0, }, { "conversation_id": "0:1:10001:20002", "peer_uid": "20002", "sender_name": "未读会话", "sender_avatar": "https://example.test/b.png", "content": "新消息", "unread_count": 0, }, ], [ { "conversation_id": "0:1:10001:20001", "peer_uid": "20001", "sender_name": "历史会话", "sender_avatar": "https://example.test/a.png", "content": "真正发生变化", "unread_count": 0, }, { "conversation_id": "0:1:10001:20002", "peer_uid": "20002", "sender_name": "未读会话", "sender_avatar": "https://example.test/b.png", "content": "新消息", "unread_count": 0, }, ], ] class _HttpClient: def __init__(self): self.get_conversations = AsyncMock(side_effect=snapshots) self.conversation_list_unsupported = False self.enter_count = 0 self.exit_count = 0 async def __aenter__(self): self.enter_count += 1 return self async def __aexit__(self, *_args): self.exit_count += 1 return False http = _HttpClient() service = DouyinImService( session=DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001), match_reply=AsyncMock(return_value=None), log_fn=AsyncMock(), account_id=9, ) service._handle_incoming = AsyncMock() controller = _Controller() with ( patch.object( service_module, "get_traffic_controller", return_value=controller, ), patch.object(service_module, "DouyinImHttpClient", return_value=http) as factory, ): # Startup indexes both previews, but only the unread conversation # is allowed to enter the reply path. await service._poll_conversations(initial=True) self.assertEqual(service._handle_incoming.await_count, 1) self.assertEqual( service._handle_incoming.await_args.args[0]["peer_uid"], "20002", ) # The first normal reconciliation sees the exact same previews; # it must not merely defer a historical-message reply explosion. service._handle_incoming.reset_mock() await service._poll_conversations() service._handle_incoming.assert_not_awaited() # A real preview transition is processed even if unread_count is # unavailable/zero on the upstream response. await service._poll_conversations() service._handle_incoming.assert_awaited_once() self.assertEqual( service._handle_incoming.await_args.args[0]["peer_uid"], "20001", ) self.assertEqual(factory.call_count, 3) self.assertEqual(http.enter_count, 3) self.assertEqual(http.exit_count, 3) self.assertEqual(controller.startup_flags, [True, False, False]) async def test_ready_is_not_blocked_by_slow_initial_unread_handler(self): events: list[str] = [] ready = asyncio.Event() handler_started = asyncio.Event() handler_cancelled = asyncio.Event() never_release = asyncio.Event() def on_ready(): events.append("ready") ready.set() async def slow_handler(_message): events.append("handler") handler_started.set() try: await never_release.wait() except asyncio.CancelledError: handler_cancelled.set() raise class _WsClient: connected = False def __init__(self, *_args, **_kwargs): self.start = AsyncMock() self.stop = AsyncMock() service = DouyinImService( session=DouyinImSession(cookies={"sessionid": "test"}, my_uid=10001), match_reply=AsyncMock(return_value=None), log_fn=AsyncMock(), account_id=901, on_ready=on_ready, ) service._verify_account_uid = AsyncMock() service._poll_conversations = AsyncMock( return_value=[ { "conversation_id": "0:1:10001:29001", "content": "startup unread", "unread_count": 1, } ] ) service._handle_incoming = AsyncMock(side_effect=slow_handler) service._reply_queue.start = AsyncMock() service._reply_queue.stop = AsyncMock() with ( patch.object(service_module, "DouyinImWsClient", _WsClient), patch.object(service_module, "ensure_frontier_ws"), patch.object(service_module.system_logger, "record"), patch( "rpa_engine.douyin_im.emoji_pack.is_fresh", return_value=True, ), ): run_task = asyncio.create_task(service.run()) try: # If initial unread were still processed inline, this wait # would time out because slow_handler never completes. await asyncio.wait_for(ready.wait(), timeout=0.3) await asyncio.wait_for(handler_started.wait(), timeout=0.3) self.assertEqual(events[:2], ["ready", "handler"]) service._poll_conversations.assert_awaited_once_with( initial=True, defer_handlers=True, ) # Account stop cancels its active deferred handler instead of # leaving work detached from the service lifecycle. await asyncio.wait_for(service.stop(), timeout=0.5) self.assertTrue(handler_cancelled.is_set()) finally: run_task.cancel() await asyncio.gather(run_task, return_exceptions=True) await service_module._shutdown_initial_unread_dispatcher() async def test_initial_unread_dispatcher_has_process_wide_concurrency_limit(self): dispatcher = service_module._InitialUnreadDispatcher(concurrency=2) active = 0 maximum_active = 0 processed: list[int] = [] two_started = asyncio.Event() release = asyncio.Event() def make_service(account_id: int): async def handle(_message): nonlocal active, maximum_active active += 1 maximum_active = max(maximum_active, active) if active == 2: two_started.set() try: await release.wait() processed.append(account_id) finally: active -= 1 return SimpleNamespace( account_id=account_id, _running=True, _handle_incoming=handle, ) services = [make_service(index + 1) for index in range(6)] try: for service in services: await dispatcher.submit(service, [{"unread_count": 1}]) await asyncio.wait_for(two_started.wait(), timeout=0.3) await asyncio.sleep(0.03) self.assertEqual(maximum_active, 2) self.assertEqual(processed, []) release.set() await asyncio.wait_for(dispatcher.join(), timeout=0.5) self.assertEqual(maximum_active, 2) self.assertCountEqual(processed, range(1, 7)) finally: await dispatcher.stop() if __name__ == "__main__": unittest.main()