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