347 lines
13 KiB
Python
347 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
|
if str(BACKEND_DIR) not in sys.path:
|
|
sys.path.insert(0, str(BACKEND_DIR))
|
|
|
|
from rpa_engine.batch_start import BatchStartQueue
|
|
from rpa_engine import batch_start as batch_start_module
|
|
|
|
|
|
class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
|
|
def _make_queue(
|
|
self,
|
|
handler,
|
|
*,
|
|
concurrency: int = 2,
|
|
timeout_seconds: float | None = None,
|
|
) -> BatchStartQueue:
|
|
queue = BatchStartQueue(
|
|
handler,
|
|
concurrency=concurrency,
|
|
timeout_seconds=timeout_seconds,
|
|
)
|
|
self.addAsyncCleanup(queue.stop)
|
|
return queue
|
|
|
|
async def _wait_for_complete(
|
|
self,
|
|
queue: BatchStartQueue,
|
|
batch_id: str,
|
|
*,
|
|
timeout: float = 0.5,
|
|
) -> dict:
|
|
async def wait() -> dict:
|
|
while True:
|
|
snapshot = await queue.get_batch(batch_id, include_items=True)
|
|
self.assertIsNotNone(snapshot)
|
|
if snapshot["complete"]:
|
|
return snapshot
|
|
await asyncio.sleep(0.001)
|
|
|
|
return await asyncio.wait_for(wait(), timeout=timeout)
|
|
|
|
async def test_concurrency_limit_is_two(self):
|
|
active = 0
|
|
maximum_active = 0
|
|
started: list[int] = []
|
|
two_started = asyncio.Event()
|
|
release = asyncio.Event()
|
|
|
|
async def handler(account_id: int) -> dict:
|
|
nonlocal active, maximum_active
|
|
active += 1
|
|
maximum_active = max(maximum_active, active)
|
|
started.append(account_id)
|
|
if len(started) == 2:
|
|
two_started.set()
|
|
try:
|
|
await release.wait()
|
|
return {"message": f"started-{account_id}", "login_mode": "im_direct"}
|
|
finally:
|
|
active -= 1
|
|
|
|
queue = self._make_queue(handler, concurrency=2)
|
|
submitted = await queue.submit(list(range(1, 9)))
|
|
|
|
await asyncio.wait_for(two_started.wait(), timeout=0.2)
|
|
await asyncio.sleep(0.02)
|
|
|
|
self.assertEqual(len(started), 2)
|
|
self.assertEqual(maximum_active, 2)
|
|
self.assertEqual(active, 2)
|
|
|
|
release.set()
|
|
completed = await self._wait_for_complete(queue, submitted["batch_id"])
|
|
|
|
self.assertEqual(completed["submitted_count"], 8)
|
|
self.assertEqual(completed["failed_count"], 0)
|
|
self.assertEqual(maximum_active, 2)
|
|
self.assertEqual(active, 0)
|
|
|
|
async def test_submit_returns_while_handler_is_blocked(self):
|
|
handler_started = asyncio.Event()
|
|
release_handler = asyncio.Event()
|
|
handler_finished = asyncio.Event()
|
|
|
|
async def handler(account_id: int) -> dict:
|
|
handler_started.set()
|
|
await release_handler.wait()
|
|
handler_finished.set()
|
|
return {"message": f"started-{account_id}"}
|
|
|
|
queue = self._make_queue(handler, concurrency=1)
|
|
|
|
submitted = await asyncio.wait_for(queue.submit([11]), timeout=0.2)
|
|
|
|
self.assertEqual(submitted["accepted_count"], 1)
|
|
self.assertEqual(submitted["queued_count"], 1)
|
|
self.assertFalse(submitted["complete"])
|
|
self.assertFalse(handler_finished.is_set())
|
|
|
|
await asyncio.wait_for(handler_started.wait(), timeout=0.2)
|
|
processing = await queue.get_batch(submitted["batch_id"])
|
|
self.assertEqual(processing["processing_count"], 1)
|
|
self.assertFalse(handler_finished.is_set())
|
|
|
|
release_handler.set()
|
|
completed = await self._wait_for_complete(queue, submitted["batch_id"])
|
|
self.assertTrue(handler_finished.is_set())
|
|
self.assertEqual(completed["submitted_count"], 1)
|
|
|
|
async def test_duplicate_ids_and_overlapping_batches_are_deduplicated(self):
|
|
first_started = asyncio.Event()
|
|
release_first = asyncio.Event()
|
|
calls: list[int] = []
|
|
|
|
async def handler(account_id: int) -> dict:
|
|
calls.append(account_id)
|
|
if account_id == 7:
|
|
first_started.set()
|
|
await release_first.wait()
|
|
return {"message": f"started-{account_id}"}
|
|
|
|
queue = self._make_queue(handler, concurrency=1)
|
|
first = await queue.submit([7, 7])
|
|
self.assertEqual(first["total_count"], 1)
|
|
self.assertEqual(first["accepted_count"], 1)
|
|
await asyncio.wait_for(first_started.wait(), timeout=0.2)
|
|
|
|
overlapping = await queue.submit([7, 8, 8])
|
|
|
|
self.assertEqual(overlapping["total_count"], 2)
|
|
self.assertEqual(overlapping["accepted_count"], 1)
|
|
self.assertEqual(overlapping["skipped_count"], 1)
|
|
by_account = {item["account_id"]: item for item in overlapping["items"]}
|
|
self.assertEqual(by_account[7]["status"], "already_queued")
|
|
self.assertEqual(by_account[8]["status"], "queued")
|
|
|
|
release_first.set()
|
|
await self._wait_for_complete(queue, first["batch_id"])
|
|
completed_overlap = await self._wait_for_complete(
|
|
queue,
|
|
overlapping["batch_id"],
|
|
)
|
|
|
|
self.assertEqual(calls, [7, 8])
|
|
self.assertEqual(completed_overlap["submitted_count"], 1)
|
|
self.assertEqual(completed_overlap["skipped_count"], 1)
|
|
|
|
async def test_failure_does_not_block_following_accounts(self):
|
|
calls: list[int] = []
|
|
|
|
async def handler(account_id: int) -> dict:
|
|
calls.append(account_id)
|
|
if account_id == 31:
|
|
raise RuntimeError("expected startup failure")
|
|
return {"message": f"started-{account_id}"}
|
|
|
|
queue = self._make_queue(handler, concurrency=1)
|
|
with patch.object(batch_start_module.logger, "exception"):
|
|
submitted = await queue.submit([31, 32, 33])
|
|
completed = await self._wait_for_complete(queue, submitted["batch_id"])
|
|
|
|
self.assertEqual(calls, [31, 32, 33])
|
|
self.assertEqual(completed["failed_count"], 1)
|
|
self.assertEqual(completed["submitted_count"], 2)
|
|
by_account = {item["account_id"]: item for item in completed["items"]}
|
|
self.assertEqual(by_account[31]["status"], "failed")
|
|
self.assertEqual(by_account[31]["message"], "expected startup failure")
|
|
self.assertEqual(by_account[32]["status"], "submitted")
|
|
self.assertEqual(by_account[33]["status"], "submitted")
|
|
|
|
async def test_two_timeouts_release_both_workers_for_following_accounts(self):
|
|
never_release = asyncio.Event()
|
|
calls: list[int] = []
|
|
|
|
async def handler(account_id: int) -> dict:
|
|
calls.append(account_id)
|
|
if account_id in (71, 72):
|
|
await never_release.wait()
|
|
return {"message": f"started-{account_id}"}
|
|
|
|
queue = self._make_queue(
|
|
handler,
|
|
concurrency=2,
|
|
timeout_seconds=0.02,
|
|
)
|
|
with patch.object(batch_start_module.logger, "warning"):
|
|
submitted = await queue.submit([71, 72, 73, 74])
|
|
completed = await self._wait_for_complete(
|
|
queue,
|
|
submitted["batch_id"],
|
|
)
|
|
|
|
self.assertEqual(calls, [71, 72, 73, 74])
|
|
self.assertEqual(completed["failed_count"], 2)
|
|
self.assertEqual(completed["submitted_count"], 2)
|
|
by_account = {item["account_id"]: item for item in completed["items"]}
|
|
self.assertIn("已跳过并继续处理后续账号", by_account[71]["message"])
|
|
self.assertIn("已跳过并继续处理后续账号", by_account[72]["message"])
|
|
self.assertEqual(by_account[73]["status"], "submitted")
|
|
self.assertEqual(by_account[74]["status"], "submitted")
|
|
|
|
async def test_handler_timeout_error_keeps_its_original_detail(self):
|
|
async def handler(_account_id: int) -> dict:
|
|
raise asyncio.TimeoutError("upstream request timed out")
|
|
|
|
queue = self._make_queue(
|
|
handler,
|
|
concurrency=1,
|
|
timeout_seconds=10,
|
|
)
|
|
with patch.object(batch_start_module.logger, "exception"):
|
|
submitted = await queue.submit([75])
|
|
completed = await self._wait_for_complete(queue, submitted["batch_id"])
|
|
|
|
item = completed["items"][0]
|
|
self.assertEqual(item["status"], "failed")
|
|
self.assertEqual(item["message"], "upstream request timed out")
|
|
|
|
async def test_failed_account_can_be_submitted_again(self):
|
|
attempts = 0
|
|
|
|
async def handler(account_id: int) -> dict:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
if attempts == 1:
|
|
raise RuntimeError("fail once")
|
|
return {"message": f"started-{account_id}"}
|
|
|
|
queue = self._make_queue(handler, concurrency=1)
|
|
with patch.object(batch_start_module.logger, "exception"):
|
|
first = await queue.submit([40])
|
|
first_completed = await self._wait_for_complete(queue, first["batch_id"])
|
|
|
|
second = await queue.submit([40])
|
|
second_completed = await self._wait_for_complete(queue, second["batch_id"])
|
|
|
|
self.assertEqual(first_completed["failed_count"], 1)
|
|
self.assertEqual(second["accepted_count"], 1)
|
|
self.assertEqual(second["skipped_count"], 0)
|
|
self.assertEqual(second_completed["submitted_count"], 1)
|
|
self.assertEqual(attempts, 2)
|
|
|
|
async def test_cancel_then_immediate_resubmit_keeps_new_job_deduplicated(self):
|
|
blocker_started = asyncio.Event()
|
|
release_blocker = asyncio.Event()
|
|
replacement_started = asyncio.Event()
|
|
release_replacement = asyncio.Event()
|
|
calls: list[int] = []
|
|
|
|
async def handler(account_id: int) -> dict:
|
|
calls.append(account_id)
|
|
if account_id == 60:
|
|
blocker_started.set()
|
|
await release_blocker.wait()
|
|
elif account_id == 61:
|
|
replacement_started.set()
|
|
await release_replacement.wait()
|
|
return {"message": f"started-{account_id}"}
|
|
|
|
queue = self._make_queue(handler, concurrency=1)
|
|
blocker = await queue.submit([60])
|
|
await asyncio.wait_for(blocker_started.wait(), timeout=0.2)
|
|
|
|
cancelled_batch = await queue.submit([61])
|
|
self.assertEqual(await queue.cancel_account(61), 1)
|
|
|
|
replacement = await queue.submit([61])
|
|
self.assertEqual(replacement["accepted_count"], 1)
|
|
|
|
# Releasing the blocker makes the worker consume the stale cancelled
|
|
# queue entry before it starts the replacement. Cleanup for that old
|
|
# token must not release the replacement's pending ownership.
|
|
release_blocker.set()
|
|
await asyncio.wait_for(replacement_started.wait(), timeout=0.2)
|
|
|
|
duplicate = await queue.submit([61])
|
|
self.assertEqual(duplicate["accepted_count"], 0)
|
|
self.assertEqual(duplicate["skipped_count"], 1)
|
|
|
|
release_replacement.set()
|
|
await self._wait_for_complete(queue, blocker["batch_id"])
|
|
cancelled = await self._wait_for_complete(
|
|
queue,
|
|
cancelled_batch["batch_id"],
|
|
)
|
|
completed = await self._wait_for_complete(queue, replacement["batch_id"])
|
|
|
|
self.assertEqual(cancelled["cancelled_count"], 1)
|
|
self.assertEqual(completed["submitted_count"], 1)
|
|
self.assertEqual(calls, [60, 61])
|
|
|
|
async def test_stop_cancels_active_and_queued_jobs_and_cleans_workers(self):
|
|
active = 0
|
|
two_started = asyncio.Event()
|
|
started: list[int] = []
|
|
never_release = asyncio.Event()
|
|
|
|
async def handler(account_id: int) -> dict:
|
|
nonlocal active
|
|
active += 1
|
|
started.append(account_id)
|
|
if len(started) == 2:
|
|
two_started.set()
|
|
try:
|
|
await never_release.wait()
|
|
return {"message": f"started-{account_id}"}
|
|
finally:
|
|
active -= 1
|
|
|
|
queue = self._make_queue(handler, concurrency=2)
|
|
submitted = await queue.submit([51, 52, 53, 54])
|
|
await asyncio.wait_for(two_started.wait(), timeout=0.2)
|
|
|
|
await asyncio.wait_for(queue.stop(), timeout=0.2)
|
|
completed = await queue.get_batch(submitted["batch_id"])
|
|
|
|
self.assertTrue(completed["complete"])
|
|
self.assertEqual(completed["cancelled_count"], 4)
|
|
self.assertEqual(active, 0)
|
|
self.assertEqual(queue._pending_jobs, {})
|
|
self.assertEqual(queue._workers, [])
|
|
self.assertTrue(queue._queue.empty())
|
|
await asyncio.wait_for(queue._queue.join(), timeout=0.1)
|
|
|
|
leaked = [
|
|
task
|
|
for task in asyncio.all_tasks()
|
|
if task is not asyncio.current_task()
|
|
and task.get_name().startswith("account-batch-start-")
|
|
and not task.done()
|
|
]
|
|
self.assertEqual(leaked, [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|