This commit is contained in:
Your Name
2026-07-30 10:06:53 +08:00
parent 8f68af1c2c
commit 3fc94c4a89
9 changed files with 275 additions and 10 deletions
+49
View File
@@ -213,6 +213,55 @@ class BatchStartApiTests(unittest.IsolatedAsyncioTestCase):
)
self.assertEqual(result["status"], "running")
async def test_batch_start_holds_no_db_connection_while_it_waits(self):
"""A queued start must not pin one of the few pooled connections.
Credential validation and readiness waiting take seconds per account.
Holding a session open across them exhausted the pool during a bulk
start, so every unrelated request waited out ``pool_timeout``.
"""
account = SimpleNamespace(
id=505,
status="starting",
qr_code_base64=None,
error_message=None,
im_session_data="saved-session",
)
events: list[str] = []
assessment = {
"login_mode": "im_direct",
"should_reset": False,
"can_skip_browser": True,
"message": "ready",
"cookie_valid": True,
"im_ready": True,
}
async def commit():
events.append("release")
async def assess(*_args, **_kwargs):
events.append("assess")
return assessment
async def start_worker(*_args, **_kwargs):
events.append("start-worker")
return True
db = SimpleNamespace(commit=AsyncMock(side_effect=commit))
with (
patch.object(main.manager, "is_running", return_value=False),
patch.object(main.manager, "start_worker", AsyncMock(side_effect=start_worker)),
patch.object(main, "_get_account_cookie_data", return_value="{}"),
patch.object(main, "assess_account_credential", AsyncMock(side_effect=assess)),
):
await main._start_account_rpa_impl(account, db, wait_for_ready=True)
# "starting" was already persisted, so the only commits here exist to
# return the connection: one before validation, one before the wait.
self.assertEqual(events, ["release", "assess", "release", "start-worker"])
async def test_batch_start_does_not_launch_interactive_browser_login(self):
account = SimpleNamespace(
id=504,
+36
View File
@@ -87,6 +87,42 @@ class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase):
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()
@@ -152,6 +152,9 @@ class CookieCredentialLockTests(unittest.IsolatedAsyncioTestCase):
[
"lock-enter",
"authorize",
# Checks the pooled connection back in before cancel/stop,
# which may wait on an in-flight start.
"commit",
"cancel",
"stop",
"write-cookie",
@@ -245,6 +248,9 @@ class CookieCredentialLockTests(unittest.IsolatedAsyncioTestCase):
[
"lock-enter",
"authorize",
# Checks the pooled connection back in before cancel/stop,
# which may wait on an in-flight start.
"commit",
"cancel",
"stop",
"clear-cookie-file",
+81
View File
@@ -389,6 +389,10 @@ class GlobalSendQueueTests(unittest.IsolatedAsyncioTestCase):
class BackgroundTrafficLimitTests(unittest.IsolatedAsyncioTestCase):
async def _wait_until(self, predicate) -> None:
while not predicate():
await asyncio.sleep(0.001)
async def test_background_slot_respects_configured_concurrency_limit(self):
with patch.dict(
os.environ,
@@ -458,6 +462,83 @@ class BackgroundTrafficLimitTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(len(tasks_seen), 3)
self.assertTrue(all(task is tasks_seen[0] for task in tasks_seen))
async def test_startup_traffic_leaves_one_slot_for_recurring_work(self):
with patch.dict(
os.environ,
{"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "3"},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
entered: list[str] = []
release = asyncio.Event()
async def request(name: str, *, startup: bool = False) -> None:
async with controller.background_slot(1, name, startup=startup):
entered.append(name)
await release.wait()
starts = [
asyncio.create_task(request(f"startup-{index}", startup=True))
for index in range(3)
]
while len(entered) < 2:
await asyncio.sleep(0)
await asyncio.sleep(0.01)
# A flood of startups may occupy at most capacity - 1 slots, so a
# recurring poll still gets in while the fleet is coming online.
self.assertEqual(len(entered), 2)
self.assertEqual(controller.background_startup_active, 2)
poll = asyncio.create_task(request("poll"))
await asyncio.wait_for(
self._wait_until(lambda: "poll" in entered),
timeout=0.5,
)
release.set()
await asyncio.wait_for(asyncio.gather(*starts, poll), timeout=0.5)
self.assertEqual(controller.background_active, 0)
self.assertEqual(controller.background_startup_active, 0)
async def test_recurring_work_is_not_deferred_indefinitely_by_startups(self):
"""Pending startup work may delay a recurring poll, never block it.
Starting several hundred accounts keeps startup requests queued for the
whole run. Yielding to that queue without a deadline left every hosted
account silent until the last account had finished coming online.
"""
with patch.dict(
os.environ,
{
"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1",
"KEFU_BACKGROUND_NORMAL_MAX_DEFER_SECONDS": "0.02",
},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
entered: list[str] = []
release = asyncio.Event()
async def poll() -> None:
async with controller.background_slot(1, "conversation poll"):
entered.append("poll")
await release.wait()
# Stands in for a batch whose startup requests never stop arriving.
controller._background_startup_clear.clear()
task = asyncio.create_task(poll())
await asyncio.wait_for(
self._wait_until(lambda: entered == ["poll"]),
timeout=1.0,
)
release.set()
await asyncio.wait_for(task, timeout=0.5)
async def test_startup_request_is_not_buried_behind_normal_backlog(self):
with patch.dict(
os.environ,