240 lines
8.2 KiB
Python
240 lines
8.2 KiB
Python
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()
|