更新
This commit is contained in:
Your Name
@@ -0,0 +1,191 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
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.session import DouyinImSession
|
||||
from rpa_engine.playwright_worker import DouyinWorker
|
||||
|
||||
|
||||
class WorkerScaleControlTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_storage_load_selects_only_cookie_column(self):
|
||||
class _ScalarResult:
|
||||
def scalar_one_or_none(self):
|
||||
return '{"cookies": [{"name": "sessionid", "value": "ok"}]}'
|
||||
|
||||
db = SimpleNamespace(
|
||||
execute=AsyncMock(return_value=_ScalarResult()),
|
||||
close=AsyncMock(),
|
||||
)
|
||||
worker = DouyinWorker(account_id=499)
|
||||
worker.get_db = AsyncMock(return_value=db)
|
||||
|
||||
storage = await worker._load_storage_state()
|
||||
|
||||
self.assertEqual(storage["cookies"][0]["value"], "ok")
|
||||
statement = db.execute.await_args.args[0]
|
||||
selected_names = [
|
||||
item.get("name") for item in statement.column_descriptions
|
||||
]
|
||||
self.assertEqual(selected_names, ["cookie_data"])
|
||||
self.assertNotIn("im_session_data", str(statement).lower())
|
||||
|
||||
async def test_direct_service_marks_worker_ready_after_initialization(self):
|
||||
worker = DouyinWorker(account_id=500, login_mode="im_direct")
|
||||
worker._refresh_follow_welcome_config = AsyncMock(
|
||||
return_value=(False, "", "sec-user")
|
||||
)
|
||||
worker.get_reply_delay = AsyncMock(return_value=None)
|
||||
fake_service = SimpleNamespace(run=AsyncMock(), stop=AsyncMock())
|
||||
|
||||
async def complete_initialization():
|
||||
service_factory.call_args.kwargs["on_ready"]()
|
||||
|
||||
fake_service.run.side_effect = complete_initialization
|
||||
session = DouyinImSession(
|
||||
cookies={"sessionid": "session"},
|
||||
my_uid=0,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"rpa_engine.playwright_worker.DouyinImService",
|
||||
return_value=fake_service,
|
||||
) as service_factory:
|
||||
await worker._run_im_direct_service(session)
|
||||
|
||||
await worker.wait_until_ready()
|
||||
worker._refresh_follow_welcome_config.assert_awaited_once_with(force=True)
|
||||
fake_service.run.assert_awaited_once_with()
|
||||
fake_service.stop.assert_awaited_once_with()
|
||||
|
||||
async def test_direct_start_failure_releases_readiness_waiter(self):
|
||||
worker = DouyinWorker(account_id=501, login_mode="im_direct")
|
||||
worker._load_storage_state = AsyncMock(return_value=None)
|
||||
worker.update_account_status = AsyncMock()
|
||||
worker.cleanup = AsyncMock()
|
||||
|
||||
await worker.start()
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "未保存 Cookie"):
|
||||
await worker.wait_until_ready()
|
||||
|
||||
if worker._task:
|
||||
await worker._task
|
||||
worker.update_account_status.assert_awaited_once_with(
|
||||
"error",
|
||||
error_msg="未保存 Cookie，无法直连 IM",
|
||||
)
|
||||
|
||||
async def test_prevalidated_start_skips_duplicate_remote_validation(self):
|
||||
worker = DouyinWorker(
|
||||
account_id=502,
|
||||
login_mode="im_direct",
|
||||
credential_prevalidated=True,
|
||||
)
|
||||
session = DouyinImSession(
|
||||
cookies={"sessionid": "session"},
|
||||
my_uid=10001,
|
||||
keys_str=json.dumps({"ec_privateKey": "private"}),
|
||||
web_protect_str=json.dumps(
|
||||
{
|
||||
"ticket": "ticket",
|
||||
"ts_sign": "sign",
|
||||
"client_cert": "certificate",
|
||||
}
|
||||
),
|
||||
)
|
||||
worker._load_user_agent = AsyncMock(return_value="test-agent")
|
||||
worker._build_im_session_from_storage = AsyncMock(return_value=session)
|
||||
worker._require_sec_user_id = AsyncMock(return_value="sec-user")
|
||||
worker._persist_im_session = AsyncMock()
|
||||
worker._run_im_direct_service = AsyncMock()
|
||||
|
||||
with patch(
|
||||
"rpa_engine.playwright_worker.validate_im_session",
|
||||
new_callable=AsyncMock,
|
||||
) as validate:
|
||||
started, reason = await worker._try_cookie_only_im_start(
|
||||
{"cookies": []}
|
||||
)
|
||||
|
||||
self.assertTrue(started)
|
||||
self.assertEqual(reason, "")
|
||||
validate.assert_not_awaited()
|
||||
worker._require_sec_user_id.assert_awaited_once()
|
||||
worker._run_im_direct_service.assert_awaited_once_with(session)
|
||||
|
||||
async def test_disabled_follow_welcome_uses_cached_lightweight_config(self):
|
||||
class _Result:
|
||||
def first(self):
|
||||
return False, "", "sec-user"
|
||||
|
||||
db = SimpleNamespace(
|
||||
execute=AsyncMock(return_value=_Result()),
|
||||
close=AsyncMock(),
|
||||
)
|
||||
worker = DouyinWorker(account_id=503)
|
||||
worker.get_db = AsyncMock(return_value=db)
|
||||
|
||||
first = await worker._refresh_follow_welcome_config()
|
||||
second = await worker._refresh_follow_welcome_config()
|
||||
|
||||
self.assertEqual(first, (False, "", "sec-user"))
|
||||
self.assertEqual(second, first)
|
||||
worker.get_db.assert_awaited_once_with()
|
||||
db.execute.assert_awaited_once()
|
||||
db.close.assert_awaited_once_with()
|
||||
|
||||
# The account PUT endpoint calls this synchronous hook so enabling the
|
||||
# feature does not wait for the disabled-account ten-minute TTL.
|
||||
worker.invalidate_follow_welcome_config()
|
||||
await worker._refresh_follow_welcome_config()
|
||||
self.assertEqual(worker.get_db.await_count, 2)
|
||||
self.assertEqual(db.execute.await_count, 2)
|
||||
|
||||
async def test_disabled_follow_tick_does_not_read_follower_log(self):
|
||||
worker = DouyinWorker(account_id=504)
|
||||
worker._im_service = SimpleNamespace(session=object())
|
||||
worker._follow_config_loaded = True
|
||||
worker._refresh_follow_welcome_config = AsyncMock(
|
||||
return_value=(False, "", "sec-user")
|
||||
)
|
||||
worker.get_db = AsyncMock()
|
||||
worker._require_sec_user_id = AsyncMock()
|
||||
|
||||
await worker.follow_welcome_tick()
|
||||
|
||||
worker._refresh_follow_welcome_config.assert_awaited_once_with()
|
||||
worker.get_db.assert_not_awaited()
|
||||
worker._require_sec_user_id.assert_not_awaited()
|
||||
|
||||
async def test_missing_cached_sec_user_id_stops_hosting(self):
|
||||
worker = DouyinWorker(account_id=505)
|
||||
worker._im_service = SimpleNamespace(session=object())
|
||||
worker._follow_config_loaded = True
|
||||
worker._refresh_follow_welcome_config = AsyncMock(
|
||||
return_value=(False, "", "")
|
||||
)
|
||||
worker._require_sec_user_id = AsyncMock(return_value="")
|
||||
|
||||
await worker.follow_welcome_tick()
|
||||
|
||||
worker._require_sec_user_id.assert_awaited_once_with(
|
||||
"托管运行中"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user