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_terminal_token_error_does_not_probe_other_payloads(self): client = self._make_client() client._request = AsyncMock( return_value={ "status_code": 500, "error_desc": "empty token", "body": {}, } ) self.assertEqual(await client.get_conversations(), []) client._request.assert_awaited_once() self.assertEqual(client._request.await_args.args[0], "POST") async def test_successful_empty_response_stops_after_first_payload(self): client = self._make_client() client._request = AsyncMock( return_value={"status_code": 0, "body": {"conversation_list": []}} ) self.assertEqual(await client.get_conversations(), []) client._request.assert_awaited_once() async def test_parameter_error_can_fall_through_to_compatible_payload(self): client = self._make_client() client._request = AsyncMock( side_effect=[ {"status_code": 400, "error_desc": "invalid parameter"}, {"status_code": 0, "body": {"conversation_list": []}}, ] ) self.assertEqual(await client.get_conversations(), []) self.assertEqual(client._request.await_count, 2) self.assertTrue( all(call.args[0] == "POST" for call in client._request.await_args_list) ) async def test_get_fallback_only_runs_after_transport_failure(self): client = self._make_client() client._request = AsyncMock( side_effect=[None, {"status_code": 0, "body": {}}] ) self.assertEqual(await client.get_conversations(), []) self.assertEqual( [call.args[0] for call in client._request.await_args_list], ["POST", "GET"], ) async def test_transport_outage_stops_after_one_post_and_get_pair(self): client = self._make_client() client._request = AsyncMock(return_value=None) with self.assertLogs("douyin_im.http", level="WARNING"): self.assertEqual(await client.get_conversations(), []) self.assertEqual(client._request.await_count, 2) self.assertEqual( [call.args[0] for call in client._request.await_args_list], ["POST", "GET"], ) 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=[]) 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.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()