更新
This commit is contained in:
Your Name
@@ -0,0 +1,239 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user