xiaobai-review/tests/test_account_switch_grants.py

from __future__ import annotations

import threading
import unittest
from datetime import datetime, timedelta, timezone
from pathlib import Path
from tempfile import TemporaryDirectory

from backend.features.accounts.security import SecretVault, token_hash
from backend.features.accounts.service import AccountService
from database import ReviewDatabase


class AccountSwitchGrantTests(unittest.TestCase):
    def setUp(self) -> None:
        self.temp = TemporaryDirectory()
        self.database = ReviewDatabase(Path(self.temp.name) / "review.db")
        self.bound_user_id = 0
        self.service = AccountService(
            database=self.database,
            vault=SecretVault(SecretVault.generate_key()),
            current_user_supplier=lambda: self.bound_user_id,
            access_supplier=lambda: self.database.user_access(self.bound_user_id) or {},
            bind_user=self._bind,
            personal_field_builder=lambda *args, **kwargs: {},
            auth_lock=threading.Lock(),
        )
        self.device_a = token_hash("device-a-token")
        self.device_b = token_hash("device-b-token")

    def tearDown(self) -> None:
        self.temp.cleanup()

    def _bind(self, user_id: int) -> None:
        self.bound_user_id = int(user_id)

    def _register(self, username: str, device_hash: str = "") -> dict:
        invite = ""
        if self.database.count_users():
            invite = self.service.generate_invite_codes(1)[0]["code"]
        return self.service.register(
            username, "Password123", device_hash or self.device_a, invite
        )

    def test_login_records_accounts_for_the_current_device_only(self) -> None:
        first = self._register("alpha_user")
        second = self._register("beta_user")
        self.service.login("alpha_user", "Password123", self.device_b)

        listed = self.service.list_device_accounts(self.device_a)
        names = [item["username"] for item in listed["accounts"]]
        self.assertEqual(names, ["beta_user", "alpha_user"])
        self.assertEqual(
            self.service.list_device_accounts(self.device_b)["accounts"][0]["username"],
            "alpha_user",
        )
        self.assertEqual(self.service.list_device_accounts("")["accounts"], [])
        self.assertEqual(first["user"]["username"], "alpha_user")
        self.assertEqual(second["user"]["username"], "beta_user")

    def test_switch_uses_device_grant_and_keeps_the_original_authorization(self) -> None:
        first = self._register("alpha_user")
        self._register("beta_user")
        switched = self.service.switch_account(self.device_a, int(first["user"]["id"]))
        self.assertEqual(switched["user"]["username"], "alpha_user")
        remaining = {
            item["username"]
            for item in self.service.list_device_accounts(self.device_a)["accounts"]
        }
        self.assertEqual(remaining, {"alpha_user", "beta_user"})

    def test_switch_without_a_valid_grant_requires_reauthentication(self) -> None:
        user = self._register("alpha_user")
        with self.assertRaisesRegex(PermissionError, "该账号需重新验证"):
            self.service.switch_account(self.device_b, int(user["user"]["id"]))
        with self.assertRaisesRegex(PermissionError, "该账号需重新验证"):
            self.service.switch_account("", int(user["user"]["id"]))

    def test_forget_only_removes_the_current_device_grant(self) -> None:
        user = self._register("alpha_user")
        self.service.login("alpha_user", "Password123", self.device_b)
        self.service.forget_account(self.device_a, int(user["user"]["id"]))
        self.service.forget_account(self.device_a, int(user["user"]["id"]))
        self.assertEqual(self.service.list_device_accounts(self.device_a)["accounts"], [])
        self.assertEqual(
            self.service.list_device_accounts(self.device_b)["accounts"][0]["username"],
            "alpha_user",
        )

    def test_logout_revokes_only_the_current_account_on_this_device(self) -> None:
        first = self._register("alpha_user")
        second = self._register("beta_user")
        self.service.revoke_current_device_grant(self.device_a, int(second["user"]["id"]))
        names = {
            item["username"]
            for item in self.service.list_device_accounts(self.device_a)["accounts"]
        }
        self.assertEqual(names, {"alpha_user"})
        switched = self.service.switch_account(self.device_a, int(first["user"]["id"]))
        self.assertEqual(switched["user"]["id"], first["user"]["id"])

    def test_password_change_revokes_grants_on_every_device(self) -> None:
        user = self._register("alpha_user")
        self.service.login("alpha_user", "Password123", self.device_b)
        self._bind(int(user["user"]["id"]))
        self.service.change_password("Password123", "Password456")
        self.assertEqual(self.service.list_device_accounts(self.device_a)["accounts"], [])
        self.assertEqual(self.service.list_device_accounts(self.device_b)["accounts"], [])
        with self.assertRaisesRegex(PermissionError, "该账号需重新验证"):
            self.service.switch_account(self.device_a, int(user["user"]["id"]))

    def test_device_keeps_at_most_five_accounts(self) -> None:
        usernames = [f"user_{index}" for index in range(6)]
        ids = [self._register(name)["user"]["id"] for name in usernames]
        listed = self.service.list_device_accounts(self.device_a)["accounts"]
        self.assertEqual(len(listed), 5)
        kept = {item["user_id"] for item in listed}
        self.assertNotIn(ids[0], kept)
        self.assertTrue(set(ids[1:]).issubset(kept))

    def test_expired_grants_are_removed_lazily(self) -> None:
        user = self._register("alpha_user")
        past = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat(timespec="seconds")
        self.database.upsert_switch_grant(
            self.device_a,
            int(user["user"]["id"]),
            past,
            past,
            past,
        )
        self.assertEqual(self.service.list_device_accounts(self.device_a)["accounts"], [])
        with self.assertRaisesRegex(PermissionError, "该账号需重新验证"):
            self.service.switch_account(self.device_a, int(user["user"]["id"]))


if __name__ == "__main__":
    unittest.main()