270 lines
9.5 KiB
Python
270 lines
9.5 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
|
|
from rpa_engine import account_profile as account_profile_module
|
|
|
|
|
|
class CookieCredentialLockTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_invalid_cookie_does_not_stop_or_lock_hosting(self):
|
|
db = SimpleNamespace()
|
|
user = SimpleNamespace(id=9, role="operator")
|
|
|
|
with (
|
|
patch.object(
|
|
main,
|
|
"validate_cookie_json",
|
|
side_effect=ValueError("invalid json"),
|
|
),
|
|
patch.object(main, "get_owned_account", new=AsyncMock()) as get_owned,
|
|
patch.object(
|
|
main.batch_start_queue,
|
|
"cancel_account",
|
|
new=AsyncMock(),
|
|
) as cancel,
|
|
patch.object(main.manager, "stop_worker", new=AsyncMock()) as stop,
|
|
patch.object(main.manager, "preparation_lock") as preparation_lock,
|
|
):
|
|
with self.assertRaises(main.HTTPException) as caught:
|
|
await main.update_account_cookie(
|
|
account_id=81,
|
|
body=main.AccountCookieUpdate(cookie_data="{bad-json"),
|
|
db=db,
|
|
user=user,
|
|
)
|
|
|
|
self.assertEqual(caught.exception.status_code, 400)
|
|
get_owned.assert_not_awaited()
|
|
cancel.assert_not_awaited()
|
|
stop.assert_not_awaited()
|
|
preparation_lock.assert_not_called()
|
|
|
|
async def test_cookie_update_stays_locked_through_stop_and_commit(self):
|
|
account = SimpleNamespace(
|
|
id=82,
|
|
cookie_data="old-cookie",
|
|
cookie_path="old-path",
|
|
cookie_updated_at=None,
|
|
updated_at=None,
|
|
)
|
|
user = SimpleNamespace(id=9, role="operator")
|
|
events: list[str] = []
|
|
lock_active = False
|
|
|
|
def require_lock(event: str) -> None:
|
|
self.assertTrue(lock_active, f"{event} ran outside preparation lock")
|
|
events.append(event)
|
|
|
|
@asynccontextmanager
|
|
async def preparation_lock(account_id: int):
|
|
nonlocal lock_active
|
|
self.assertEqual(account_id, 82)
|
|
lock_active = True
|
|
events.append("lock-enter")
|
|
try:
|
|
yield
|
|
finally:
|
|
events.append("lock-exit")
|
|
lock_active = False
|
|
|
|
async def get_owned(db, selected_user, account_id, *, write=False):
|
|
require_lock("authorize")
|
|
self.assertIs(db, fake_db)
|
|
self.assertIs(selected_user, user)
|
|
self.assertEqual(account_id, 82)
|
|
self.assertTrue(write)
|
|
return account
|
|
|
|
async def cancel_account(account_id: int):
|
|
require_lock("cancel")
|
|
self.assertEqual(account_id, 82)
|
|
return 0
|
|
|
|
async def stop_worker(account_id: int):
|
|
require_lock("stop")
|
|
self.assertEqual(account_id, 82)
|
|
return True
|
|
|
|
async def execute(_statement):
|
|
require_lock("clear-profile")
|
|
return None
|
|
|
|
async def commit():
|
|
require_lock("commit")
|
|
|
|
async def refresh(selected_account):
|
|
require_lock("refresh")
|
|
self.assertIs(selected_account, account)
|
|
|
|
async def apply_profile(db, selected_account, cookie_data):
|
|
require_lock("sync-profile")
|
|
self.assertIs(db, fake_db)
|
|
self.assertIs(selected_account, account)
|
|
self.assertIn("sessionid", cookie_data)
|
|
|
|
fake_db = SimpleNamespace(
|
|
execute=AsyncMock(side_effect=execute),
|
|
commit=AsyncMock(side_effect=commit),
|
|
refresh=AsyncMock(side_effect=refresh),
|
|
)
|
|
|
|
def write_cookie(account_id: int, cookie_data: str) -> str:
|
|
require_lock("write-cookie")
|
|
self.assertEqual(account_id, 82)
|
|
self.assertIn("sessionid", cookie_data)
|
|
return "new-cookie-path"
|
|
|
|
with (
|
|
patch.object(main, "validate_cookie_json", return_value={"cookies": [{"name": "sessionid", "value": "new"}]}),
|
|
patch.object(main, "get_owned_account", new=AsyncMock(side_effect=get_owned)) as get_owned_mock,
|
|
patch.object(main.manager, "preparation_lock", side_effect=preparation_lock),
|
|
patch.object(main.batch_start_queue, "cancel_account", new=AsyncMock(side_effect=cancel_account)) as cancel,
|
|
patch.object(main.manager, "stop_worker", new=AsyncMock(side_effect=stop_worker)) as stop,
|
|
patch.object(main, "write_cookie_file", side_effect=write_cookie),
|
|
patch.object(account_profile_module, "apply_douyin_profile", new=AsyncMock(side_effect=apply_profile)),
|
|
patch.object(main, "_build_cookie_response", new=AsyncMock(return_value={"account_id": 82})) as build_response,
|
|
):
|
|
response = await main.update_account_cookie(
|
|
account_id=82,
|
|
body=main.AccountCookieUpdate(cookie_data="valid-cookie"),
|
|
db=fake_db,
|
|
user=user,
|
|
)
|
|
|
|
self.assertEqual(response, {"account_id": 82})
|
|
self.assertEqual(
|
|
events,
|
|
[
|
|
"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",
|
|
"clear-profile",
|
|
"sync-profile",
|
|
"commit",
|
|
"refresh",
|
|
"lock-exit",
|
|
],
|
|
)
|
|
get_owned_mock.assert_awaited_once()
|
|
cancel.assert_awaited_once_with(82)
|
|
stop.assert_awaited_once_with(82)
|
|
self.assertEqual(account.cookie_path, "new-cookie-path")
|
|
build_response.assert_awaited_once()
|
|
|
|
async def test_cookie_delete_stays_locked_through_stop_and_commit(self):
|
|
account = SimpleNamespace(
|
|
id=83,
|
|
cookie_data="old-cookie",
|
|
cookie_path="old-path",
|
|
cookie_updated_at=object(),
|
|
im_session_data="old-im-session",
|
|
updated_at=None,
|
|
)
|
|
user = SimpleNamespace(id=9, role="operator")
|
|
events: list[str] = []
|
|
lock_active = False
|
|
|
|
def require_lock(event: str) -> None:
|
|
self.assertTrue(lock_active, f"{event} ran outside preparation lock")
|
|
events.append(event)
|
|
|
|
@asynccontextmanager
|
|
async def preparation_lock(account_id: int):
|
|
nonlocal lock_active
|
|
self.assertEqual(account_id, 83)
|
|
lock_active = True
|
|
events.append("lock-enter")
|
|
try:
|
|
yield
|
|
finally:
|
|
events.append("lock-exit")
|
|
lock_active = False
|
|
|
|
async def get_owned(_db, _user, account_id, *, write=False):
|
|
require_lock("authorize")
|
|
self.assertEqual(account_id, 83)
|
|
self.assertTrue(write)
|
|
return account
|
|
|
|
async def cancel_account(_account_id: int):
|
|
require_lock("cancel")
|
|
return 0
|
|
|
|
async def stop_worker(_account_id: int):
|
|
require_lock("stop")
|
|
return True
|
|
|
|
async def execute(_statement):
|
|
require_lock("clear-profile")
|
|
|
|
async def commit():
|
|
require_lock("commit")
|
|
|
|
fake_db = SimpleNamespace(
|
|
execute=AsyncMock(side_effect=execute),
|
|
commit=AsyncMock(side_effect=commit),
|
|
)
|
|
|
|
def clear_cookie(account_id: int) -> None:
|
|
require_lock("clear-cookie-file")
|
|
self.assertEqual(account_id, 83)
|
|
|
|
with (
|
|
patch.object(main, "get_owned_account", new=AsyncMock(side_effect=get_owned)),
|
|
patch.object(main.manager, "preparation_lock", side_effect=preparation_lock),
|
|
patch.object(main.batch_start_queue, "cancel_account", new=AsyncMock(side_effect=cancel_account)),
|
|
patch.object(main.manager, "stop_worker", new=AsyncMock(side_effect=stop_worker)),
|
|
patch.object(main, "clear_cookie_file", side_effect=clear_cookie),
|
|
):
|
|
response = await main.delete_account_cookie(
|
|
account_id=83,
|
|
db=fake_db,
|
|
user=user,
|
|
)
|
|
|
|
self.assertEqual(response["message"], "Cookie cleared successfully.")
|
|
self.assertEqual(
|
|
events,
|
|
[
|
|
"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",
|
|
"clear-profile",
|
|
"commit",
|
|
"lock-exit",
|
|
],
|
|
)
|
|
self.assertIsNone(account.cookie_data)
|
|
self.assertIsNone(account.cookie_path)
|
|
self.assertIsNone(account.cookie_updated_at)
|
|
self.assertIsNone(account.im_session_data)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|