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.http_client import DouyinImHttpClient from rpa_engine.douyin_im.session import DouyinImSession 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) 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_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", ) self.assertFalse(sent) submit.assert_awaited_once() queued_factory.assert_called_once_with(client.session, account_id=88) queued_client.send_text_message.assert_awaited_once_with( "0:1:10001:20002", "queued hello", conversation_short_id="short-before-send", _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") class WorkerLifecycleTests(unittest.IsolatedAsyncioTestCase): 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()