285 lines
10 KiB
Python
285 lines
10 KiB
Python
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()
|