更新
This commit is contained in:
Your Name
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
@@ -32,6 +33,219 @@ def _fake_db(rows):
|
||||
|
||||
|
||||
class BatchStartApiTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_shutdown_continues_when_batch_queue_cleanup_times_out(self):
|
||||
original_workers = main.manager.workers
|
||||
original_flush_task = main._system_log_flush_task
|
||||
main.manager.workers = {}
|
||||
main._system_log_flush_task = None
|
||||
stop_started = asyncio.Event()
|
||||
|
||||
async def blocked_batch_stop():
|
||||
stop_started.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
try:
|
||||
with (
|
||||
patch.dict(
|
||||
os.environ,
|
||||
{"KEFU_BATCH_STOP_TIMEOUT_SECONDS": "1"},
|
||||
),
|
||||
patch.object(
|
||||
main.batch_start_queue,
|
||||
"stop",
|
||||
AsyncMock(side_effect=blocked_batch_stop),
|
||||
),
|
||||
patch(
|
||||
"rpa_engine.douyin_im.traffic_control.shutdown_traffic_controller",
|
||||
AsyncMock(),
|
||||
) as stop_traffic,
|
||||
):
|
||||
await asyncio.wait_for(main.shutdown(), timeout=2.0)
|
||||
|
||||
self.assertTrue(stop_started.is_set())
|
||||
stop_traffic.assert_awaited_once_with()
|
||||
finally:
|
||||
main.manager.workers = original_workers
|
||||
main._system_log_flush_task = original_flush_task
|
||||
|
||||
async def test_shutdown_stops_many_accounts_with_bounded_parallelism(self):
|
||||
active = 0
|
||||
maximum_active = 0
|
||||
stopped: list[int] = []
|
||||
original_workers = main.manager.workers
|
||||
original_flush_task = main._system_log_flush_task
|
||||
main.manager.workers = {
|
||||
account_id: SimpleNamespace(is_running=True)
|
||||
for account_id in range(601, 613)
|
||||
}
|
||||
main._system_log_flush_task = None
|
||||
|
||||
async def stop_worker(account_id: int):
|
||||
nonlocal active, maximum_active
|
||||
active += 1
|
||||
maximum_active = max(maximum_active, active)
|
||||
try:
|
||||
await asyncio.sleep(0.005)
|
||||
stopped.append(account_id)
|
||||
main.manager.workers.pop(account_id, None)
|
||||
return True
|
||||
finally:
|
||||
active -= 1
|
||||
|
||||
try:
|
||||
with (
|
||||
patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"KEFU_SHUTDOWN_CONCURRENCY": "3",
|
||||
"KEFU_SHUTDOWN_TIMEOUT_SECONDS": "5",
|
||||
},
|
||||
),
|
||||
patch.object(main.batch_start_queue, "stop", AsyncMock()),
|
||||
patch.object(main.manager, "stop_worker", AsyncMock(side_effect=stop_worker)),
|
||||
patch(
|
||||
"rpa_engine.douyin_im.traffic_control.shutdown_traffic_controller",
|
||||
AsyncMock(),
|
||||
),
|
||||
):
|
||||
await main.shutdown()
|
||||
|
||||
self.assertEqual(len(stopped), 12)
|
||||
self.assertEqual(maximum_active, 3)
|
||||
finally:
|
||||
main.manager.workers = original_workers
|
||||
main._system_log_flush_task = original_flush_task
|
||||
|
||||
async def test_worker_manager_waits_for_full_ready_and_reuses_validation(self):
|
||||
worker = SimpleNamespace(
|
||||
is_running=True,
|
||||
start=AsyncMock(),
|
||||
wait_until_ready=AsyncMock(),
|
||||
)
|
||||
manager = main.WorkerManager()
|
||||
|
||||
with patch.object(main, "DouyinWorker", return_value=worker) as worker_factory:
|
||||
started = await manager.start_worker(
|
||||
501,
|
||||
login_mode="im_direct",
|
||||
wait_until_ready=True,
|
||||
credential_prevalidated=True,
|
||||
)
|
||||
|
||||
self.assertTrue(started)
|
||||
worker_factory.assert_called_once_with(
|
||||
501,
|
||||
login_mode="im_direct",
|
||||
credential_prevalidated=True,
|
||||
)
|
||||
worker.start.assert_awaited_once_with()
|
||||
worker.wait_until_ready.assert_awaited_once_with()
|
||||
self.assertIs(manager.workers[501], worker)
|
||||
|
||||
async def test_cancelled_ready_wait_stops_and_removes_detached_worker(self):
|
||||
wait_started = asyncio.Event()
|
||||
waiting = asyncio.Event()
|
||||
|
||||
async def wait_forever():
|
||||
wait_started.set()
|
||||
await waiting.wait()
|
||||
|
||||
worker = SimpleNamespace(
|
||||
is_running=True,
|
||||
start=AsyncMock(),
|
||||
wait_until_ready=AsyncMock(side_effect=wait_forever),
|
||||
stop=AsyncMock(),
|
||||
)
|
||||
manager = main.WorkerManager()
|
||||
|
||||
with patch.object(main, "DouyinWorker", return_value=worker):
|
||||
task = asyncio.create_task(
|
||||
manager.start_worker(
|
||||
502,
|
||||
login_mode="im_direct",
|
||||
wait_until_ready=True,
|
||||
credential_prevalidated=True,
|
||||
)
|
||||
)
|
||||
await asyncio.wait_for(wait_started.wait(), timeout=0.2)
|
||||
task.cancel()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
worker.stop.assert_awaited_once_with()
|
||||
self.assertNotIn(502, manager.workers)
|
||||
|
||||
async def test_batch_start_waits_for_ready_and_skips_duplicate_validation(self):
|
||||
account = SimpleNamespace(
|
||||
id=503,
|
||||
status="offline",
|
||||
qr_code_base64=None,
|
||||
error_message=None,
|
||||
im_session_data="saved-session",
|
||||
)
|
||||
db = SimpleNamespace(commit=AsyncMock())
|
||||
assessment = {
|
||||
"login_mode": "im_direct",
|
||||
"should_reset": False,
|
||||
"can_skip_browser": True,
|
||||
"message": "ready",
|
||||
"cookie_valid": True,
|
||||
"im_ready": True,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(main.manager, "is_running", return_value=False),
|
||||
patch.object(main.manager, "start_worker", AsyncMock(return_value=True)) as start,
|
||||
patch.object(main, "_get_account_cookie_data", return_value="{}"),
|
||||
patch.object(main, "assess_account_credential", AsyncMock(return_value=assessment)),
|
||||
):
|
||||
result = await main._start_account_rpa_impl(
|
||||
account,
|
||||
db,
|
||||
wait_for_ready=True,
|
||||
)
|
||||
|
||||
start.assert_awaited_once_with(
|
||||
503,
|
||||
login_mode="im_direct",
|
||||
wait_until_ready=True,
|
||||
credential_prevalidated=True,
|
||||
)
|
||||
self.assertEqual(result["status"], "running")
|
||||
|
||||
async def test_batch_start_does_not_launch_interactive_browser_login(self):
|
||||
account = SimpleNamespace(
|
||||
id=504,
|
||||
status="offline",
|
||||
qr_code_base64=None,
|
||||
error_message=None,
|
||||
im_session_data=None,
|
||||
)
|
||||
db = SimpleNamespace(commit=AsyncMock())
|
||||
assessment = {
|
||||
"login_mode": "browser",
|
||||
"should_reset": False,
|
||||
"can_skip_browser": False,
|
||||
"message": "login required",
|
||||
"cookie_valid": False,
|
||||
"im_ready": False,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(main.manager, "is_running", return_value=False),
|
||||
patch.object(main.manager, "start_worker", AsyncMock()) as start,
|
||||
patch.object(main, "_get_account_cookie_data", return_value=None),
|
||||
patch.object(main, "assess_account_credential", AsyncMock(return_value=assessment)),
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, "批量启动"):
|
||||
await main._start_account_rpa_impl(
|
||||
account,
|
||||
db,
|
||||
wait_for_ready=True,
|
||||
)
|
||||
|
||||
start.assert_not_awaited()
|
||||
|
||||
async def test_start_all_uses_lightweight_select_and_submits_once(self):
|
||||
db = _fake_db([(1, False), (2, True), (3, False), (4, False)])
|
||||
user = SimpleNamespace(id=9, role="admin")
|
||||
|
||||
Reference in New Issue
Block a user