更新

2026-07-23 17:56:25 +08:00
parent a05dae8412
commit 4970d8f8d3
4262 changed files with 735221 additions and 0 deletions
@@ -0,0 +1,284 @@
+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()