test/test_local_sqlite_split.py

from __future__ import annotations

import os
import sqlite3
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from services.storage.database_storage import DatabaseStorageBackend
from services.storage.factory import (
    _migrate_local_sqlite_users_once,
    create_storage_backend,
    create_user_storage_backend,
)


def _table_names(path: Path) -> set[str]:
    connection = sqlite3.connect(path)
    try:
        return {
            str(row[0])
            for row in connection.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
        }
    finally:
        connection.close()


class LocalSqliteSplitTests(unittest.TestCase):
    def test_explicit_default_accounts_database_url_still_enables_split(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            data_dir = Path(tmp_dir)
            database_url = f"sqlite:///{data_dir / 'accounts.db'}"
            with patch.dict(
                os.environ,
                {
                    "STORAGE_BACKEND": "sqlite",
                    "DATABASE_URL": database_url,
                    "CHATGPT2API_SPLIT_LOCAL_DATABASES": "true",
                    "CHATGPT2API_SKIP_STORAGE_MIGRATION": "true",
                },
                clear=False,
            ):
                os.environ.pop("SHARED_USER_DATABASE_URL", None)
                account_backend = create_storage_backend(data_dir)
                user_backend = create_user_storage_backend(data_dir, account_backend)
            try:
                self.assertEqual(account_backend.table_scope, "accounts")
                self.assertEqual(user_backend.table_scope, "users")
                self.assertTrue((data_dir / "users.db").exists())
            finally:
                account_backend.engine.dispose()
                user_backend.engine.dispose()

    def test_shared_user_migration_reads_from_split_users_database(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            data_dir = Path(tmp_dir)
            legacy = DatabaseStorageBackend(f"sqlite:///{data_dir / 'accounts.db'}")
            legacy.save_auth_keys([{"id": "shared-user", "username": "共享中文用户"}])
            legacy.engine.dispose()
            shared_url = f"sqlite:///{data_dir / 'shared-users.db'}"

            with patch.dict(
                os.environ,
                {
                    "STORAGE_BACKEND": "sqlite",
                    "SHARED_USER_DATABASE_URL": shared_url,
                    "CHATGPT2API_SPLIT_LOCAL_DATABASES": "true",
                    "CHATGPT2API_AUTO_MIGRATE_USER_STORAGE": "true",
                    "CHATGPT2API_SKIP_STORAGE_MIGRATION": "true",
                },
                clear=False,
            ):
                os.environ.pop("DATABASE_URL", None)
                account_backend = create_storage_backend(data_dir)
                shared_backend = create_user_storage_backend(data_dir, account_backend)
            try:
                self.assertEqual(shared_backend.load_auth_keys()[0]["username"], "共享中文用户")
            finally:
                account_backend.engine.dispose()
                shared_backend.engine.dispose()

    def test_migration_preserves_unicode_and_is_idempotent(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            data_dir = Path(tmp_dir)
            accounts_path = data_dir / "accounts.db"
            legacy = DatabaseStorageBackend(f"sqlite:///{accounts_path}")
            try:
                legacy.save_accounts(
                    [
                        {
                            "access_token": "account-token-1",
                            "email": "pool@example.com",
                            "owner_id": "user-cn",
                            "status": "active",
                        }
                    ]
                )
                legacy.save_auth_keys(
                    [
                        {
                            "id": "user-cn",
                            "username": "听风用户",
                            "nickname": "中文昵称不会乱码",
                            "role": "normal",
                            "enabled": True,
                        }
                    ]
                )
                legacy.save_shop_state(
                    {
                        "codes": [
                            {
                                "id": "code-cn",
                                "code_hash": "hash-cn",
                                "batch_id": "批次一",
                                "status": "active",
                                "note": "中文卡密",
                            }
                        ],
                        "ledger": [
                            {
                                "id": "ledger-cn",
                                "time": "2026-08-24T12:00:00+08:00",
                                "user_id": "user-cn",
                                "note": "签到赠送",
                            }
                        ],
                    }
                )
                legacy.save_named_state("announcements", {"items": [{"title": "中文公告"}]})
            finally:
                legacy.engine.dispose()

            with patch.dict(
                os.environ,
                {
                    "STORAGE_BACKEND": "sqlite",
                    "CHATGPT2API_SPLIT_LOCAL_DATABASES": "true",
                    "CHATGPT2API_SKIP_STORAGE_MIGRATION": "true",
                },
                clear=False,
            ):
                os.environ.pop("DATABASE_URL", None)
                os.environ.pop("SHARED_USER_DATABASE_URL", None)
                _migrate_local_sqlite_users_once(data_dir)
                _migrate_local_sqlite_users_once(data_dir)

                account_backend = create_storage_backend(data_dir)
                user_backend = create_user_storage_backend(data_dir, account_backend)

            try:
                self.assertEqual(account_backend.table_scope, "accounts")
                self.assertEqual(user_backend.table_scope, "users")
                self.assertEqual(account_backend.load_accounts()[0]["email"], "pool@example.com")

                users = user_backend.load_auth_keys()
                self.assertEqual(users[0]["username"], "听风用户")
                self.assertEqual(users[0]["nickname"], "中文昵称不会乱码")
                shop_state = user_backend.load_shop_state()
                self.assertEqual(shop_state["codes"][0]["note"], "中文卡密")
                self.assertEqual(shop_state["codes"][0]["batch_id"], "批次一")
                self.assertEqual(shop_state["ledger"][0]["note"], "签到赠送")
                self.assertEqual(
                    user_backend.load_named_state("announcements")["items"][0]["title"],
                    "中文公告",
                )

                self.assertIn("accounts", _table_names(accounts_path))
                self.assertTrue(
                    {"auth_keys", "shop_state", "redeem_codes", "shop_ledger"}.isdisjoint(
                        _table_names(accounts_path)
                    )
                )
                self.assertNotIn("accounts", _table_names(data_dir / "users.db"))
                self.assertTrue(
                    {"auth_keys", "shop_state", "redeem_codes", "shop_ledger"}.issubset(
                        _table_names(data_dir / "users.db")
                    )
                )
                self.assertTrue((data_dir / "accounts.before-user-split.db").exists())
                self.assertTrue((data_dir / ".local_user_database_split_migrated").exists())
                self.assertEqual(account_backend.health_check()["status"], "healthy")
                self.assertEqual(user_backend.health_check()["status"], "healthy")
            finally:
                account_backend.engine.dispose()
                user_backend.engine.dispose()

    def test_account_database_failure_does_not_break_user_database(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            data_dir = Path(tmp_dir)
            legacy = DatabaseStorageBackend(f"sqlite:///{data_dir / 'accounts.db'}")
            legacy.save_accounts([{"access_token": "account-token"}])
            legacy.save_auth_keys([{"id": "user-1", "username": "隔离用户"}])
            legacy.engine.dispose()

            with patch.dict(
                os.environ,
                {
                    "CHATGPT2API_SPLIT_LOCAL_DATABASES": "true",
                    "CHATGPT2API_SKIP_STORAGE_MIGRATION": "true",
                },
                clear=False,
            ):
                _migrate_local_sqlite_users_once(data_dir)

            connection = sqlite3.connect(data_dir / "accounts.db")
            try:
                connection.execute("DROP TABLE accounts")
                connection.commit()
            finally:
                connection.close()

            users = DatabaseStorageBackend(f"sqlite:///{data_dir / 'users.db'}", table_scope="users")
            try:
                self.assertEqual(users.load_auth_keys()[0]["username"], "隔离用户")
                self.assertEqual(users.health_check()["status"], "healthy")
            finally:
                users.engine.dispose()


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