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()