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_default_concurrency_admits_more_than_two_accounts(self): with patch.dict(os.environ, {}, clear=False): os.environ.pop("KEFU_BATCH_START_CONCURRENCY", None) queue = BatchStartQueue(lambda _account_id: None) self.assertEqual(queue.concurrency, batch_start_module.DEFAULT_CONCURRENCY) self.assertGreaterEqual(queue.concurrency, 4) async def test_dead_worker_is_replaced_so_width_never_shrinks(self): """One crashed worker must not permanently narrow the queue. Width used to be restored only when every worker had exited, so a single unexpected worker death left later batches crawling through the survivors until the process restarted. """ async def handler(account_id: int) -> dict: return {"message": f"started-{account_id}"} queue = self._make_queue(handler, concurrency=3) first = await queue.submit([91]) await self._wait_for_complete(queue, first["batch_id"]) self.assertEqual(len(queue._workers), 3) casualty = queue._workers[0] casualty.cancel() await asyncio.gather(casualty, return_exceptions=True) second = await queue.submit([92]) await self._wait_for_complete(queue, second["batch_id"]) self.assertEqual(len(queue._workers), 3) self.assertNotIn(casualty, queue._workers) self.assertTrue(all(not task.done() for task in queue._workers)) names = [task.get_name() for task in queue._workers] self.assertEqual(len(set(names)), 3) 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()