更新

backend/tests/test_cookie_credential_lock.py

@@ -0,0 +1,263 @@
+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",
+                "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",
+                "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()