from __future__ import annotations import os import sys import unittest from contextlib import asynccontextmanager from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, 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)) import main class _RowsResult: def __init__(self, rows): self._rows = list(rows) def all(self): return list(self._rows) def _fake_db(rows): return SimpleNamespace(execute=AsyncMock(return_value=_RowsResult(rows))) class BatchStartApiTests(unittest.IsolatedAsyncioTestCase): 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") submitted_response = { "batch_id": "all-batch", "accepted_count": 2, "complete": False, } submit = AsyncMock(return_value=submitted_response) with ( patch.object(main.batch_start_queue, "submit", submit), patch.object( main.manager, "is_running", side_effect=lambda account_id: account_id == 3, ) as is_running, patch.object(main, "_build_account_response") as build_response, ): response = await main.submit_account_start_batch( body=main.BatchStartRequest(all_accounts=True), db=db, user=user, ) self.assertEqual(response, submitted_response) db.execute.assert_awaited_once() submit.assert_awaited_once_with( [1, 4], owner_id=9, metadata={ "requested_count": 4, "accessible_count": 4, "skipped_running_count": 1, "skipped_disabled_count": 1, }, ) self.assertEqual(is_running.call_count, 3) build_response.assert_not_called() statement = db.execute.await_args.args[0] selected_names = [entry.get("name") for entry in statement.column_descriptions] self.assertEqual(selected_names, ["id", "quota_disabled"]) self.assertNotIn("cookie_data", str(statement).lower()) self.assertNotIn("im_session_data", str(statement).lower()) async def test_selected_ids_are_deduplicated_scoped_and_filtered(self): db = _fake_db([(5, False), (3, False)]) user = SimpleNamespace(id=42, role="operator") submit = AsyncMock( return_value={ "batch_id": "selected-batch", "accepted_count": 2, "complete": False, } ) with ( patch.object(main.batch_start_queue, "submit", submit), patch.object(main.manager, "is_running", return_value=False), ): await main.submit_account_start_batch( body=main.BatchStartRequest( account_ids=[5, 5, 3, -1, 999], all_accounts=False, ), db=db, user=user, ) submit.assert_awaited_once_with( [5, 3], owner_id=42, metadata={ "requested_count": 3, "accessible_count": 2, "skipped_running_count": 0, "skipped_disabled_count": 0, }, ) statement = db.execute.await_args.args[0] compiled_params = statement.compile().params parameter_values = list(compiled_params.values()) self.assertIn(42, parameter_values) self.assertIn([5, 3, 999], parameter_values) sql = str(statement).lower() self.assertIn("owner_id", sql) self.assertIn("accounts.id in", sql) async def test_metadata_counts_inaccessible_running_and_disabled_accounts(self): db = _fake_db([(101, False), (102, True), (103, False)]) user = SimpleNamespace(id=7, role="operator") submit = AsyncMock(return_value={"batch_id": "metadata-batch"}) with ( patch.object(main.batch_start_queue, "submit", submit), patch.object( main.manager, "is_running", side_effect=lambda account_id: account_id == 103, ), ): await main.submit_account_start_batch( body=main.BatchStartRequest(account_ids=[101, 102, 103, 104]), db=db, user=user, ) submit.assert_awaited_once_with( [101], owner_id=7, metadata={ "requested_count": 4, "accessible_count": 3, "skipped_running_count": 1, "skipped_disabled_count": 1, }, ) async def test_batch_lookup_is_isolated_by_owner(self): async def get_batch(batch_id, owner_id, include_items=False): if batch_id == "owned-batch" and owner_id == 11: return { "batch_id": batch_id, "accepted_count": 2, "complete": True, } return None lookup = AsyncMock(side_effect=get_batch) with patch.object(main.batch_start_queue, "get_batch", lookup): owned = await main.get_account_start_batch( batch_id="owned-batch", user=SimpleNamespace(id=11, role="operator"), ) with self.assertRaises(main.HTTPException) as caught: await main.get_account_start_batch( batch_id="owned-batch", user=SimpleNamespace(id=12, role="operator"), ) self.assertEqual(owned["batch_id"], "owned-batch") self.assertEqual(caught.exception.status_code, 404) self.assertEqual( lookup.await_args_list[0].kwargs, {"owner_id": 11, "include_items": False}, ) self.assertEqual( lookup.await_args_list[1].kwargs, {"owner_id": 12, "include_items": False}, ) async def test_single_account_start_cancels_queued_batch_job_first(self): account = SimpleNamespace(id=77) db = SimpleNamespace() user = SimpleNamespace(id=5, role="operator") events: list[str] = [] async def cancel_account(account_id: int): events.append("cancel") return 1 @asynccontextmanager async def preparation_lock(account_id: int): self.assertEqual(account_id, 77) events.append("lock-enter") try: yield finally: events.append("lock-exit") async def start_impl(selected_account, selected_db, requested_login_mode): events.append("start") self.assertIs(selected_account, account) self.assertIs(selected_db, db) self.assertEqual(requested_login_mode, "im_direct") return {"status": "starting", "login_mode": "im_direct"} cancel = AsyncMock(side_effect=cancel_account) get_owned = AsyncMock(return_value=account) start = AsyncMock(side_effect=start_impl) with ( patch.object(main, "get_owned_account", get_owned), patch.object(main.batch_start_queue, "cancel_account", cancel), patch.object(main.manager, "preparation_lock", side_effect=preparation_lock), patch.object(main, "_start_account_rpa_impl", start), ): response = await main.start_account_rpa( account_id=77, body=main.StartAccountRequest(login_mode="im_direct"), db=db, user=user, ) self.assertEqual(response["status"], "starting") get_owned.assert_awaited_once_with(db, user, 77, write=True) cancel.assert_awaited_once_with(77) start.assert_awaited_once_with(account, db, "im_direct") self.assertEqual(events, ["cancel", "lock-enter", "start", "lock-exit"]) if __name__ == "__main__": unittest.main()