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, MagicMock, 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 import http_client as http_client_module from rpa_engine.douyin_im.auth import DouyinAuth from rpa_engine.douyin_im.frontier import ensure_frontier_ws from rpa_engine.douyin_im.http_client import DouyinImHttpClient from rpa_engine.douyin_im.session import DouyinImSession from rpa_engine.egress_channels import EgressChannel from rpa_engine import playwright_worker as playwright_worker_module from rpa_engine.playwright_worker import DouyinWorker class SendTextMessageEntryTests(unittest.IsolatedAsyncioTestCase): def _make_client(self, account_id: int = 41) -> DouyinImHttpClient: session = DouyinImSession( cookies={"sessionid": "test-session"}, my_uid=10001, conv_meta={"existing": {"ticket": "old-ticket"}}, ) return DouyinImHttpClient(session, account_id=account_id) def test_query_user_does_not_replace_existing_im_uid(self): client = self._make_client() auth = SimpleNamespace( source_ip="", get_uid=MagicMock(return_value=938334054809296), ) resolved = client._resolve_authoritative_uid(auth) self.assertEqual(resolved, 10001) self.assertEqual(client.session.my_uid, 10001) auth.get_uid.assert_not_called() def test_verified_uid_does_not_replace_the_runtime_device_id(self): # device_id 是 query/user 返回的设备注册号，my_uid 是账号 UID。 # 拿 my_uid 顶替 device_id 会让 frontier 订阅到另一个地址：握手照样 # 成功，却永远收不到这个账号的私信。 session = DouyinImSession( cookies={"sessionid": "test-session"}, my_uid=2609567359568155, device_id="7678285795559818786", web_id="7678286623234475535", uid_verified=True, ) auth = DouyinAuth.from_im_session(session) self.assertEqual(auth.device_id, "7678285795559818786") self.assertEqual(session.device_id, "7678285795559818786") def test_captured_tokenless_browser_frontier_url_is_kept(self): url = ( "wss://frontier100-normal.zijieapi.com/ws/v2?aid=6383&" "device_platform=web&fpid=9&device_id=7678285795559818786&" "access_key=0123456789abcdef0123456789abcdef&version_code=fws_1.0.0" ) session = DouyinImSession( cookies={"sessionid": "test-session"}, ws_urls=[url], my_uid=2609567359568155, device_id="7678285795559818786", ) resolved = ensure_frontier_ws(session) self.assertEqual(resolved, url) self.assertEqual(session.ws_urls, [url]) def test_non_im_frontier_product_is_rejected(self): url = ( "wss://frontier100-normal.zijieapi.com/ws/v2?aid=6383&" "device_platform=web&fpid=971&device_id=7678285795559818786&" "access_key=0123456789abcdef0123456789abcdef" ) session = DouyinImSession( cookies={"sessionid": "test-session"}, ws_urls=[url], my_uid=2609567359568155, device_id="7678285795559818786", uid_verified=True, ) resolved = ensure_frontier_ws(session) self.assertIn("frontier-im.douyin.com", resolved) self.assertIn("fpid=9", resolved) self.assertIn("device_id=7678285795559818786", resolved) self.assertNotIn("fpid=971", resolved) async def test_public_entry_submits_the_whole_send_to_outbound_queue(self): client = self._make_client(account_id=73) submit = AsyncMock(return_value=True) with patch( "rpa_engine.douyin_im.traffic_control.submit_outbound", submit, ): sent = await client.send_text_message( "0:1:10001:20002", "hello", conversation_short_id="short-1", ) self.assertTrue(sent) submit.assert_awaited_once() account_id, operation = submit.await_args.args[:2] self.assertEqual(account_id, 73) self.assertTrue(callable(operation)) self.assertIn("20002", submit.await_args.kwargs["description"]) async def test_bypass_path_does_not_submit_again(self): client = self._make_client() client._resolve_authoritative_uid = MagicMock(return_value=0) submit = AsyncMock(return_value=True) with patch( "rpa_engine.douyin_im.traffic_control.submit_outbound", submit, ): sent = await client.send_text_message( "0:1:10001:20002", "hello", _bypass_global_queue=True, ) self.assertFalse(sent) submit.assert_not_awaited() client._resolve_authoritative_uid.assert_called_once() async def test_queued_client_result_state_is_copied_to_calling_client(self): client = self._make_client(account_id=88) queued_meta = { "0:1:10001:20002": { "ticket": "fresh-ticket", "conversation_short_id": "fresh-short-id", } } queued_client = SimpleNamespace( send_text_message=AsyncMock(return_value=False), last_send_meta=queued_meta, last_error="credential expired", last_send_needs_refresh=True, last_send_channel_retryable=False, last_request_debug="response status=401", ) queued_context = MagicMock() queued_context.__aenter__ = AsyncMock(return_value=queued_client) queued_context.__aexit__ = AsyncMock(return_value=None) queued_factory = MagicMock(return_value=queued_context) async def execute_submission(account_id, operation, description=""): self.assertEqual(account_id, 88) self.assertIn("20002", description) return await operation() submit = AsyncMock(side_effect=execute_submission) with ( patch( "rpa_engine.douyin_im.traffic_control.submit_outbound", submit, ), patch.object( http_client_module, "DouyinImHttpClient", queued_factory, ), ): sent = await client.send_text_message( "0:1:10001:20002", "queued hello", conversation_short_id="short-before-send", expected_peer_uid="20002", ) self.assertFalse(sent) submit.assert_awaited_once() queued_factory.assert_called_once_with(client.session, account_id=88) # 收件人期望必须原样传给真正写出去的那个 client：排队调度层不能把它吃掉 queued_client.send_text_message.assert_awaited_once_with( "0:1:10001:20002", "queued hello", conversation_short_id="short-before-send", expected_peer_uid="20002", _bypass_global_queue=True, ) self.assertEqual(client.last_send_meta, queued_meta) self.assertIsNot(client.last_send_meta, queued_meta) self.assertEqual(client.last_error, "credential expired") self.assertTrue(client.last_send_needs_refresh) self.assertEqual(client.last_request_debug, "response status=401") async def test_retryable_network_failure_switches_channels_serially(self): client = self._make_client(account_id=89) client.session.egress_auto_attempts = 2 routes = [ EgressChannel("198.51.100.10", "10.0.0.10", "eth0", True), EgressChannel("198.51.100.11", "10.0.0.11", "eth0:1", False), ] first = SimpleNamespace( send_text_message=AsyncMock(return_value=False), last_send_meta={}, last_error="connect timeout", last_send_needs_refresh=False, last_send_channel_retryable=True, last_request_debug="first route", ) second = SimpleNamespace( send_text_message=AsyncMock(return_value=True), last_send_meta={"conv": {"ticket": "ok"}}, last_error="", last_send_needs_refresh=False, last_send_channel_retryable=False, last_request_debug="second route", ) def context_for(value): context = MagicMock() context.__aenter__ = AsyncMock(return_value=value) context.__aexit__ = AsyncMock(return_value=None) return context queued_factory = MagicMock(side_effect=[context_for(first), context_for(second)]) async def execute_submission(account_id, operation, description=""): self.assertEqual(account_id, 89) return await operation() with ( patch( "rpa_engine.douyin_im.traffic_control.submit_outbound", AsyncMock(side_effect=execute_submission), ), patch( "rpa_engine.douyin_im.http_client.resolve_send_channels", AsyncMock(return_value=routes), ), patch.object(http_client_module, "DouyinImHttpClient", queued_factory), patch.object(http_client_module.system_logger, "record"), ): sent = await client.send_text_message("0:1:10001:20002", "hello") self.assertTrue(sent) self.assertEqual(queued_factory.call_count, 2) self.assertEqual(queued_factory.call_args_list[0].kwargs["source_ip"], "10.0.0.10") self.assertEqual(queued_factory.call_args_list[1].kwargs["source_ip"], "10.0.0.11") first.send_text_message.assert_awaited_once() second.send_text_message.assert_awaited_once() self.assertEqual(client.last_request_debug, "second route") async def test_kick_never_switches_public_channels(self): client = self._make_client(account_id=90) client.session.egress_auto_attempts = 2 routes = [ EgressChannel("198.51.100.10", "10.0.0.10", "eth0", True), EgressChannel("198.51.100.11", "10.0.0.11", "eth1", False), ] kicked = SimpleNamespace( send_text_message=AsyncMock(return_value=False), last_send_meta={}, last_error="decision=KICK", last_send_needs_refresh=False, last_send_channel_retryable=False, last_request_debug="terminal kick", ) context = MagicMock() context.__aenter__ = AsyncMock(return_value=kicked) context.__aexit__ = AsyncMock(return_value=None) queued_factory = MagicMock(return_value=context) async def execute_submission(account_id, operation, description=""): return await operation() with ( patch( "rpa_engine.douyin_im.traffic_control.submit_outbound", AsyncMock(side_effect=execute_submission), ), patch( "rpa_engine.douyin_im.http_client.resolve_send_channels", AsyncMock(return_value=routes), ), patch.object(http_client_module, "DouyinImHttpClient", queued_factory), ): sent = await client.send_text_message("0:1:10001:20002", "hello") self.assertFalse(sent) queued_factory.assert_called_once() kicked.send_text_message.assert_awaited_once() class WorkerLifecycleTests(unittest.IsolatedAsyncioTestCase): async def test_browser_launch_uses_selected_source_proxy(self): launch = AsyncMock(return_value="browser") pw = SimpleNamespace(chromium=SimpleNamespace(launch=launch)) source_proxy = AsyncMock(return_value={"server": "http://127.0.0.1:43210"}) with ( patch.object( playwright_worker_module, "ensure_browser_display", AsyncMock(), ), patch.object( playwright_worker_module, "playwright_proxy_for_source", source_proxy, ), patch.object(playwright_worker_module, "playwright_proxy") as global_proxy, ): browser = await playwright_worker_module._launch_chromium( pw, ["--no-sandbox"], headless=True, source_ip="10.0.0.6", ) self.assertEqual(browser, "browser") source_proxy.assert_awaited_once_with("10.0.0.6") global_proxy.assert_not_called() launch.assert_awaited_once_with( headless=True, args=["--no-sandbox"], proxy={"server": "http://127.0.0.1:43210"}, ) async def test_virtual_display_starts_before_playwright_driver(self): ensure_display = AsyncMock() async def start_driver(): ensure_display.assert_awaited_once_with(False) return "playwright-driver" manager = SimpleNamespace(start=AsyncMock(side_effect=start_driver)) with ( patch.object( playwright_worker_module, "resolve_headless", return_value=False, ), patch.object( playwright_worker_module, "ensure_browser_display", ensure_display, ), patch.object( playwright_worker_module, "async_playwright", return_value=manager, ), ): driver, headless = await playwright_worker_module._start_playwright_for_browser() self.assertEqual(driver, "playwright-driver") self.assertFalse(headless) manager.start.assert_awaited_once_with() async def test_visible_login_prompt_overrides_stale_sessionid(self): worker = DouyinWorker(account_id=917, login_mode="browser") worker._has_visible_login_prompt = AsyncMock(return_value=True) worker.check_homepage_login_status = AsyncMock(return_value=True) worker.check_logged_in_by_cookie = AsyncMock(return_value=True) logged_in = await worker._verify_login_state() self.assertFalse(logged_in) worker.check_homepage_login_status.assert_not_awaited() worker.check_logged_in_by_cookie.assert_not_awaited() async def test_homepage_message_entry_is_not_login_evidence(self): worker = DouyinWorker(account_id=918, login_mode="browser") worker._has_visible_login_prompt = AsyncMock(return_value=False) worker.page = SimpleNamespace(query_selector=AsyncMock(return_value=None)) logged_in = await worker.check_homepage_login_status() self.assertFalse(logged_in) worker.page.query_selector.assert_awaited_once() async def test_start_saves_task_and_stop_waits_until_it_is_done(self): worker = DouyinWorker(account_id=919, login_mode="im_direct") loop_started = asyncio.Event() loop_finished = asyncio.Event() never_finishes_without_cancellation = asyncio.Event() async def controlled_run_loop() -> None: loop_started.set() try: await never_finishes_without_cancellation.wait() finally: loop_finished.set() worker._run_loop = controlled_run_loop worker.update_account_status = AsyncMock() worker.cleanup = AsyncMock() await worker.start() task = worker._task self.assertIsInstance(task, asyncio.Task) self.assertEqual(task.get_name(), "douyin-worker-919") self.assertFalse(task.done()) await asyncio.wait_for(loop_started.wait(), timeout=0.2) # Starting an already-running worker must retain the same owned task. await worker.start() self.assertIs(worker._task, task) await asyncio.wait_for(worker.stop(), timeout=0.5) self.assertTrue(task.done()) self.assertTrue(task.cancelled()) self.assertTrue(loop_finished.is_set()) self.assertIsNone(worker._task) self.assertTrue(worker.stopping) self.assertFalse(worker.is_running) worker.update_account_status.assert_awaited_once_with("offline") worker.cleanup.assert_awaited_once_with() leaked = [ running for running in asyncio.all_tasks() if running is not asyncio.current_task() and running.get_name() == "douyin-worker-919" and not running.done() ] self.assertEqual(leaked, []) async def test_browser_flow_closes_resources_before_releasing_slot_on_error(self): worker = DouyinWorker(account_id=920, login_mode="browser") events: list[str] = [] # A truthy browser resource makes the flow's finally block responsible # for cleanup even though credential preparation fails midway. worker.page = object() async def failing_prepare(storage_state, cookie_info): events.append("prepare") raise RuntimeError("credential harvest failed") async def close_browser_only(): events.append("close") worker.page = None @asynccontextmanager async def browser_slot(account_id, description): self.assertEqual(account_id, 920) self.assertEqual(description, "browser credential harvest") events.append("acquire") try: yield finally: events.append("release") controller = SimpleNamespace(browser_slot=browser_slot) worker._prepare_browser_im_flow = AsyncMock(side_effect=failing_prepare) worker._close_browser_only = AsyncMock(side_effect=close_browser_only) with patch( "rpa_engine.playwright_worker.get_traffic_controller", return_value=controller, ): with self.assertRaisesRegex(RuntimeError, "credential harvest failed"): await worker._run_browser_im_flow( storage_state={"cookies": []}, cookie_info={"has_sessionid": True}, ) self.assertEqual(events, ["acquire", "prepare", "close", "release"]) worker._prepare_browser_im_flow.assert_awaited_once_with( {"cookies": []}, {"has_sessionid": True}, ) worker._close_browser_only.assert_awaited_once_with() async def test_close_browser_only_closes_every_resource_and_clears_references(self): worker = DouyinWorker(account_id=921, login_mode="browser") events: list[str] = [] async def record(name: str, *, fail: bool = False) -> None: events.append(name) if fail: raise RuntimeError(f"{name} close failed") async def close_page() -> None: await record("page") async def close_context() -> None: await record("context", fail=True) async def close_browser() -> None: await record("browser") async def stop_playwright() -> None: await record("playwright") page = SimpleNamespace( close=AsyncMock(side_effect=close_page), ) context = SimpleNamespace( # One failed close must not keep the remaining processes alive. close=AsyncMock(side_effect=close_context), ) browser = SimpleNamespace( close=AsyncMock(side_effect=close_browser), ) playwright = SimpleNamespace( stop=AsyncMock(side_effect=stop_playwright), ) worker.page = page worker.context = context worker.browser = browser worker.playwright = playwright await worker._close_browser_only() self.assertEqual(events, ["page", "context", "browser", "playwright"]) page.close.assert_awaited_once_with() context.close.assert_awaited_once_with() browser.close.assert_awaited_once_with() playwright.stop.assert_awaited_once_with() self.assertIsNone(worker.page) self.assertIsNone(worker.context) self.assertIsNone(worker.browser) self.assertIsNone(worker.playwright) if __name__ == "__main__": unittest.main()