"""Tests for upstream key uniquification and sync cleanup.""" import json from datetime import datetime, timezone import pytest from sqlalchemy import create_engine, text from sqlalchemy.orm import sessionmaker from sqlalchemy.pool import StaticPool from app.database import Base from app.models.upstream import Upstream from app.models.website import Website # noqa: F401 — registers table for FK refs from app.models.upstream_key import UpstreamGeneratedKey @pytest.fixture() def db_session(): engine = create_engine( "sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool, ) Base.metadata.create_all(bind=engine) TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) db = TestingSessionLocal() try: yield db finally: db.close() Base.metadata.drop_all(bind=engine) def test_duplicate_cleanup_keeps_latest_only(): """同一 upstream_id + group_id + key_name 的多条记录只保留最新一条。 使用独立 engine + 全 raw SQL,模拟迁移前的数据库状态。 """ engine = create_engine("sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool) with engine.begin() as conn: conn.execute(text(""" CREATE TABLE upstreams ( id INTEGER PRIMARY KEY AUTOINCREMENT, name VARCHAR(255) NOT NULL, base_url VARCHAR(512) NOT NULL, api_prefix VARCHAR(64) DEFAULT '', auth_type VARCHAR(32), auth_config_json TEXT DEFAULT '{}', groups_endpoint VARCHAR(256), rate_endpoint VARCHAR(256), enabled BOOLEAN DEFAULT 1, check_interval_seconds INTEGER DEFAULT 600, timeout_seconds INTEGER DEFAULT 30, last_status VARCHAR(32) DEFAULT 'unknown', last_checked_at DATETIME, last_error TEXT, consecutive_failures INTEGER DEFAULT 0, balance FLOAT, balance_updated_at DATETIME, balance_endpoint VARCHAR(256) DEFAULT '', balance_response_path VARCHAR(256) DEFAULT '', balance_divisor FLOAT DEFAULT 1.0, updated_at DATETIME, created_at DATETIME ) """)) conn.execute(text(""" CREATE TABLE upstream_generated_keys ( id INTEGER PRIMARY KEY AUTOINCREMENT, upstream_id INTEGER NOT NULL, group_id VARCHAR(255) NOT NULL, group_name VARCHAR(255) DEFAULT '', key_id VARCHAR(255), key_name VARCHAR(255) NOT NULL, key_value TEXT NOT NULL, masked_key VARCHAR(255) DEFAULT '', raw_json TEXT DEFAULT '{}', status VARCHAR(32) DEFAULT 'created', error TEXT, imported_website_id INTEGER, imported_account_id VARCHAR(255), imported_at DATETIME, created_at DATETIME, updated_at DATETIME ) """)) TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) db = TestingSessionLocal() try: now = datetime.now(timezone.utc) db.execute(text(""" INSERT INTO upstreams (name, base_url, api_prefix, auth_type, auth_config_json, groups_endpoint, rate_endpoint, created_at, updated_at) VALUES (:n, :b, :p, :a, :j, :g, :r, :ca, :ua) """), {"n": "Test", "b": "http://local", "p": "/api/v1", "a": "bearer", "j": "{}", "g": "/groups", "r": "/rates", "ca": now, "ua": now}) db.commit() uid = db.execute(text("SELECT id FROM upstreams LIMIT 1")).scalar() # 插入 3 条重复记录 for kv, ca in [ ("old-key", datetime(2025, 1, 1, tzinfo=timezone.utc)), ("middle-key", datetime(2025, 6, 1, tzinfo=timezone.utc)), ("newest-key", datetime(2025, 12, 1, tzinfo=timezone.utc)), ]: db.execute(text(""" INSERT INTO upstream_generated_keys (upstream_id, group_id, group_name, key_name, key_value, masked_key, raw_json, status, created_at, updated_at) VALUES (:uid, :gid, :gn, :kn, :kv, :mk, :rj, :st, :ca, :ca) """), { "uid": uid, "gid": "vip", "gn": "VIP", "kn": "SmartUp-Test-VIP", "kv": kv, "mk": "", "rj": "{}", "st": "created", "ca": ca, }) db.commit() # 清理:同一组合只保留最新一条(id 最大) db.execute(text(""" DELETE FROM upstream_generated_keys WHERE id NOT IN ( SELECT MAX(id) FROM upstream_generated_keys GROUP BY upstream_id, group_id, key_name ) """)) db.commit() remaining = db.execute(text("SELECT key_value FROM upstream_generated_keys")).fetchall() assert len(remaining) == 1, f"expected 1 after cleanup, got {len(remaining)}" assert remaining[0][0] == "newest-key" finally: db.close() def test_migration_backfills_managed_prefix_and_deduplicates(): """迁移逻辑应回填历史 SmartUp 记录的 managed_prefix 并清理重复。 使用独立 engine(不创建唯一约束),模拟迁移前状态。 """ from sqlalchemy import text as _text engine = create_engine("sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool) with engine.begin() as conn: conn.execute(_text(""" CREATE TABLE upstreams ( id INTEGER PRIMARY KEY AUTOINCREMENT, name VARCHAR(255) NOT NULL, base_url VARCHAR(512) NOT NULL, api_prefix VARCHAR(64) DEFAULT '', auth_type VARCHAR(32), auth_config_json TEXT DEFAULT '{}', groups_endpoint VARCHAR(256), rate_endpoint VARCHAR(256), enabled BOOLEAN DEFAULT 1, check_interval_seconds INTEGER DEFAULT 600, timeout_seconds INTEGER DEFAULT 30, last_status VARCHAR(32) DEFAULT 'unknown', last_checked_at DATETIME, last_error TEXT, consecutive_failures INTEGER DEFAULT 0, balance FLOAT, balance_updated_at DATETIME, balance_endpoint VARCHAR(256) DEFAULT '', balance_response_path VARCHAR(256) DEFAULT '', balance_divisor FLOAT DEFAULT 1.0, updated_at DATETIME, created_at DATETIME ) """)) conn.execute(_text(""" CREATE TABLE upstream_generated_keys ( id INTEGER PRIMARY KEY AUTOINCREMENT, upstream_id INTEGER NOT NULL, group_id VARCHAR(255) NOT NULL, group_name VARCHAR(255) DEFAULT '', key_id VARCHAR(255), key_name VARCHAR(255) NOT NULL, key_value TEXT NOT NULL, masked_key VARCHAR(255) DEFAULT '', raw_json TEXT DEFAULT '{}', managed_prefix VARCHAR(64), status VARCHAR(32) DEFAULT 'created', error TEXT, imported_website_id INTEGER, imported_account_id VARCHAR(255), imported_at DATETIME, created_at DATETIME, updated_at DATETIME ) """)) TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) db = TestingSessionLocal() try: now = datetime.now(timezone.utc) db.execute(_text(""" INSERT INTO upstreams (name, base_url, api_prefix, auth_type, auth_config_json, groups_endpoint, rate_endpoint, created_at, updated_at) VALUES (:n, :b, :p, :a, :j, :g, :r, :ca, :ua) """), {"n": "Old", "b": "http://local", "p": "/api/v1", "a": "bearer", "j": "{}", "g": "/groups", "r": "/rates", "ca": now, "ua": now}) db.commit() uid = db.execute(_text("SELECT id FROM upstreams LIMIT 1")).scalar() # 插入两条重复记录(无 managed_prefix,key_name 以 SmartUp 开头) for kv in ("sk-old", "sk-new"): db.execute(_text(""" INSERT INTO upstream_generated_keys (upstream_id, group_id, group_name, key_name, key_value, masked_key, raw_json, status, created_at, updated_at) VALUES (:uid, :gid, :gn, :kn, :kv, '', '{}', 'created', :ca, :ca) """), {"uid": uid, "gid": "vip", "gn": "VIP", "kn": "SmartUp-Old-vip", "kv": kv, "ca": now}) db.commit() # 执行迁移逻辑(与 database.py 中的 SQL 一致) conn = db.connection() conn.execute(_text( "UPDATE upstream_generated_keys SET managed_prefix = 'SmartUp' " "WHERE managed_prefix IS NULL AND key_name LIKE 'SmartUp-%'" )) to_delete = conn.execute(_text(""" SELECT id FROM upstream_generated_keys WHERE managed_prefix IS NOT NULL AND id NOT IN ( SELECT MAX(id) FROM upstream_generated_keys WHERE managed_prefix IS NOT NULL GROUP BY upstream_id, group_id, managed_prefix ) """)).fetchall() for (row_id,) in to_delete: conn.execute(_text("DELETE FROM upstream_generated_keys WHERE id = :id"), {"id": row_id}) db.commit() remaining = db.execute(_text("SELECT key_value, managed_prefix FROM upstream_generated_keys")).fetchall() assert len(remaining) == 1, f"expected 1 after migration, got {len(remaining)}" assert remaining[0][0] == "sk-new" # 保留最新一条 assert remaining[0][1] == "SmartUp" # 已回填 finally: db.close() def test_ensure_group_key_reuses_old_record(db_session, monkeypatch): """_ensure_group_key 应复用 managed_prefix IS NULL 的旧记录,不新建。""" from app.routers.upstreams import _ensure_group_key from app.models.upstream_key import UpstreamGeneratedKey from app.services.upstream_client import UpstreamClient from app.schemas.upstream import GenerateKeysByGroupsRequest upstream = Upstream(name="Test", base_url="http://local", api_prefix="/api/v1", auth_type="bearer", auth_config_json="{}", groups_endpoint="/groups", rate_endpoint="/rates") db_session.add(upstream) db_session.commit() db_session.refresh(upstream) # 插入一条旧记录(无 managed_prefix) db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="vip", group_name="VIP", key_name="SmartUp-Test-vip", key_value="sk-old", managed_prefix=None, key_id="remote-999", )) db_session.commit() # 构造 mock client class MockClient: def find_smartup_group_key(self, gid, name, prefix): return None def create_api_key(self, name, group_id, **kw): return {"id": "new-remote", "key": "sk-new-value", "masked_key": "sk-****-lue"} group = {"id": "vip", "name": "VIP", "rate_multiplier": 1} body = GenerateKeysByGroupsRequest( group_ids=["vip"], name_prefix="SmartUp", quota=0, endpoint="/keys", ) result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body) assert result.status == "created" rows = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.group_id == "vip", ).all() assert len(rows) == 1, f"expected 1 record, got {len(rows)}" assert rows[0].managed_prefix == "SmartUp" assert rows[0].key_value == "sk-new-value" def test_ensure_group_key_backfills_plaintext_from_remote_existing_key(db_session): """远端已存在的 SmartUp Key 如果列表接口返回明文,应补写到本地 key_value。""" from app.routers.upstreams import _ensure_group_key from app.models.upstream_key import UpstreamGeneratedKey from app.schemas.upstream import GenerateKeysByGroupsRequest from app.services.upstream_client import mask_secret upstream = Upstream(name="Test", base_url="http://local", api_prefix="/api/v1", auth_type="bearer", auth_config_json="{}", groups_endpoint="/groups", rate_endpoint="/rates") db_session.add(upstream) db_session.commit() db_session.refresh(upstream) db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="vip", group_name="VIP", key_name="SmartUp-Test-vip", key_value="", masked_key="sk-old-masked", key_id="remote-123", managed_prefix="SmartUp", )) db_session.commit() class MockClient: def find_smartup_group_key(self, gid, name, prefix): return { "id": "remote-123", "name": "SmartUp-Test-vip", "key": "sk-remote-plain-value-1234567890abcdef", "masked_key": "sk-re************cdef", } def create_api_key(self, *args, **kwargs): raise AssertionError("create_api_key should not be called when remote key exists") group = {"id": "vip", "name": "VIP", "rate_multiplier": 1} body = GenerateKeysByGroupsRequest( group_ids=["vip"], name_prefix="SmartUp", quota=0, endpoint="/keys", ) result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body) assert result.status == "exists" assert result.has_key_value is True row = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.group_id == "vip", ).one() assert row.key_value == "sk-remote-plain-value-1234567890abcdef" assert row.masked_key == mask_secret(row.key_value) assert row.status == "exists" def test_sync_removes_remote_key_when_list_empty(db_session, monkeypatch): """同步函数在远端返回空列表时应删除本地 key_id 对应的记录。""" from app.services import scheduler as sched_mod from app.models.upstream_key import UpstreamGeneratedKey from app.services.upstream_client import UpstreamClient upstream = Upstream(name="Test", base_url="http://local", api_prefix="/api/v1", auth_type="bearer", auth_config_json="{}", groups_endpoint="/groups", rate_endpoint="/rates") db_session.add(upstream) db_session.commit() db_session.refresh(upstream) db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="vip", group_name="VIP", key_name="SmartUp-Test-vip", key_value="sk-vip", managed_prefix="SmartUp", key_id="remote-key-id", )) db_session.commit() # mock list_api_keys 返回空列表(查询成功但无 Key) monkeypatch.setattr(UpstreamClient, "list_api_keys", lambda self, **kw: []) monkeypatch.setattr(UpstreamClient, "login", lambda self: None) monkeypatch.setattr(UpstreamClient, "close", lambda self: None) monkeypatch.setattr(UpstreamClient, "__enter__", lambda self: self) monkeypatch.setattr(UpstreamClient, "__exit__", lambda self, *a: None) # 让 _sync_upstream_keys 使用 db_session 的 bind 引擎 monkeypatch.setattr(sched_mod, "SessionLocal", lambda: db_session) # 阻止 finally 中的 db.close() 影响测试会话 original_close = db_session.close monkeypatch.setattr(db_session, "close", lambda: None) snapshot = { "upstream_id": upstream.id, "groups": {"vip": {"group_id": "vip", "rate": "1"}}, "captured_at": datetime.now(timezone.utc).isoformat(), } captured_at = datetime.now(timezone.utc) sched_mod._sync_upstream_keys(upstream.id, snapshot, captured_at) monkeypatch.setattr(db_session, "close", original_close) remaining = db_session.query(UpstreamGeneratedKey).all() assert len(remaining) == 0, f"expected 0 after sync with empty remote, got {len(remaining)}" def test_sync_marks_imported_key_orphaned_when_remote_key_missing(db_session, monkeypatch): """已导入账号管理的 Key 远端消失时保留本地行,避免丢失目标账号关联。""" from app.services import scheduler as sched_mod from app.models.upstream_key import UpstreamGeneratedKey from app.services.upstream_client import UpstreamClient website = Website( name="Target", site_type="sub2api", base_url="http://target.local", api_prefix="/api/v1/admin", auth_type="api_key", auth_config_json="{}", groups_endpoint="/groups", group_update_endpoint="/groups/{id}", ) upstream = Upstream(name="Test", base_url="http://local", api_prefix="/api/v1", auth_type="bearer", auth_config_json="{}", groups_endpoint="/groups", rate_endpoint="/rates") db_session.add_all([website, upstream]) db_session.commit() db_session.refresh(website) db_session.refresh(upstream) db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="vip", group_name="VIP", key_name="SmartUp-Test-vip", key_value="sk-vip", managed_prefix="SmartUp", key_id="remote-key-id", imported_website_id=website.id, imported_account_id="account-101", )) db_session.commit() monkeypatch.setattr(UpstreamClient, "list_api_keys", lambda self, **kw: []) monkeypatch.setattr(UpstreamClient, "login", lambda self: None) monkeypatch.setattr(UpstreamClient, "close", lambda self: None) monkeypatch.setattr(UpstreamClient, "__enter__", lambda self: self) monkeypatch.setattr(UpstreamClient, "__exit__", lambda self, *a: None) monkeypatch.setattr(sched_mod, "SessionLocal", lambda: db_session) original_close = db_session.close monkeypatch.setattr(db_session, "close", lambda: None) snapshot = { "upstream_id": upstream.id, "groups": {"vip": {"group_id": "vip", "rate": "1"}}, "captured_at": datetime.now(timezone.utc).isoformat(), } captured_at = datetime.now(timezone.utc) sched_mod._sync_upstream_keys(upstream.id, snapshot, captured_at) monkeypatch.setattr(db_session, "close", original_close) remaining = db_session.query(UpstreamGeneratedKey).all() assert len(remaining) == 1 row = remaining[0] assert row.status == "orphaned" assert row.imported_website_id == website.id assert row.imported_account_id == "account-101" assert row.error == "远端 Key 已不存在" def test_migration_function_integration(monkeypatch): """直接调用 _migrate_upstream_generated_keys() 验证列新增和索引创建。""" from app.database import _migrate_upstream_generated_keys, engine as real_engine from sqlalchemy import text as _text # 使用独立 engine,避免影响真实数据库 test_engine = create_engine("sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool) monkeypatch.setattr("app.database.engine", test_engine) # 建表(不含 managed_prefix 列,模拟旧版 schema) with test_engine.begin() as conn: conn.execute(_text(""" CREATE TABLE upstream_generated_keys ( id INTEGER PRIMARY KEY AUTOINCREMENT, upstream_id INTEGER NOT NULL, group_id VARCHAR(255) NOT NULL, group_name VARCHAR(255) DEFAULT '', key_id VARCHAR(255), key_name VARCHAR(255) NOT NULL, key_value TEXT NOT NULL, masked_key VARCHAR(255) DEFAULT '', raw_json TEXT DEFAULT '{}', status VARCHAR(32) DEFAULT 'created', error TEXT, imported_website_id INTEGER, imported_account_id VARCHAR(255), imported_at DATETIME, created_at DATETIME ) """)) conn.execute(_text(""" INSERT INTO upstream_generated_keys (upstream_id, group_id, group_name, key_name, key_value, masked_key, raw_json, status, created_at) VALUES (1, 'vip', 'VIP', 'SmartUp-Old-vip', 'sk-val', '', '{}', 'created', datetime('now')) """)) # 调用迁移函数入口 _migrate_upstream_generated_keys() # 验证 managed_prefix 列已存在且被填充 inspector = __import__('sqlalchemy', fromlist=['']).inspect(test_engine) cols = {c["name"] for c in inspector.get_columns("upstream_generated_keys")} assert "managed_prefix" in cols, "managed_prefix column should exist after migration" with test_engine.connect() as conn: row = conn.execute(_text("SELECT managed_prefix, key_value FROM upstream_generated_keys LIMIT 1")).fetchone() assert row[0] == "SmartUp", f"expected SmartUp, got {row[0]}" assert row[1] == "sk-val" # 验证唯一索引已创建 indexes = inspector.get_indexes("upstream_generated_keys") index_names = {ix["name"] for ix in indexes} assert "uq_upstream_group_managed" in index_names, "partial unique index should exist" monkeypatch.undo() def test_create_twice_only_one_record(db_session): """同一上游同一分组连续调用两次 ensure,本地只保留一条记录。""" from app.models.upstream_key import UpstreamGeneratedKey upstream = Upstream(name="Test", base_url="http://local", api_prefix="/api/v1", auth_type="bearer", auth_config_json="{}", groups_endpoint="/groups", rate_endpoint="/rates") db_session.add(upstream) db_session.commit() db_session.refresh(upstream) # 模拟第一次创建 db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="vip", group_name="VIP", key_name="SmartUp-Test-VIP", key_value="sk-first", status="created", )) db_session.commit() # 模拟第二次调用 upsert(用同一个 key_name 且 status=exists) existing = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.group_id == "vip", UpstreamGeneratedKey.key_name == "SmartUp-Test-VIP", ).first() if existing: existing.status = "exists" existing.updated_at = datetime.now(timezone.utc) else: db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="vip", group_name="VIP", key_name="SmartUp-Test-VIP", key_value="sk-second", status="exists", )) db_session.commit() rows = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.group_id == "vip", ).all() assert len(rows) == 1 assert rows[0].status == "exists" assert rows[0].key_value == "sk-first" # 更新的是原记录,不是新建 def test_sync_removes_gone_group(db_session): """分组不在最新快照中时,本地对应 Key 记录应被删除。""" upstream = Upstream(name="Test", base_url="http://local", api_prefix="/api/v1", auth_type="bearer", auth_config_json="{}", groups_endpoint="/groups", rate_endpoint="/rates") db_session.add(upstream) db_session.commit() db_session.refresh(upstream) db_session.add_all([ UpstreamGeneratedKey( upstream_id=upstream.id, group_id="vip", group_name="VIP", key_name="SmartUp-Test-VIP", key_value="sk-vip", ), UpstreamGeneratedKey( upstream_id=upstream.id, group_id="free", group_name="Free", key_name="SmartUp-Test-Free", key_value="sk-free", ), ]) db_session.commit() # 快照中只有 vip,没有 free active_group_ids = {"vip"} for row in db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id).all(): if row.group_id not in active_group_ids: db_session.delete(row) db_session.commit() remaining = db_session.query(UpstreamGeneratedKey).all() assert len(remaining) == 1 assert remaining[0].group_id == "vip" def test_sync_removes_deleted_remote_key(db_session): """远端 Key 被删除后,本地对应记录应被删除。""" from app.models.upstream_key import UpstreamGeneratedKey upstream = Upstream(name="Test", base_url="http://local", api_prefix="/api/v1", auth_type="bearer", auth_config_json="{}", groups_endpoint="/groups", rate_endpoint="/rates") db_session.add(upstream) db_session.commit() db_session.refresh(upstream) db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="vip", group_name="VIP", key_name="SmartUp-Test-VIP", key_value="sk-vip", key_id="remote-123", )) db_session.commit() # 模拟远端返回的活跃 key_ids 中没有 remote-123 remote_key_ids = {"remote-456", "remote-789"} for row in db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id).all(): if row.key_id and row.key_id not in remote_key_ids: db_session.delete(row) db_session.commit() remaining = db_session.query(UpstreamGeneratedKey).all() assert len(remaining) == 0 def test_new_api_create_token_fetches_plaintext_key(monkeypatch): """New-API 创建 token 后需按 id 再取一次明文 key。""" from app.services.upstream_client import UpstreamClient client = UpstreamClient( base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config={"cookie_string": "session=abc", "new_api_user": "7"}, ) created_bodies = [] def fake_request(method, path, body=None, auth=True): if method == "GET" and path == "/api/status": return {"success": True, "data": {"quota_per_unit": 500000}} if method == "POST" and path == "/api/token/": created_bodies.append(body) return {"success": True, "message": ""} if method == "POST" and path == "/api/token/123/key": return {"success": True, "data": {"key": "new-api-plain-key"}} raise AssertionError(f"unexpected request {method} {path}") monkeypatch.setattr(client, "_request", fake_request) monkeypatch.setattr( client, "_list_new_api_tokens", lambda search="", group_id=None: [{"id": 123, "name": search, "group": group_id, "key": "new-****-key"}], ) result = client.create_api_key( "SmartUp-1-VIP-vip", "vip", quota=2, expires_in_days=3, endpoint="/api/token", ) assert result["id"] == "123" assert result["key"] == "new-api-plain-key" assert created_bodies[0]["group"] == "vip" assert created_bodies[0]["remain_quota"] == 1000000 assert created_bodies[0]["unlimited_quota"] is False assert created_bodies[0]["expired_time"] > 0 def test_new_api_list_tokens_uses_full_list_and_fetches_plaintext_when_search_misses(): """New-API search 可能不匹配前缀;应拉完整 token 列表并按 id 回填明文。""" from app.services.upstream_client import UpstreamClient client = UpstreamClient( base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config={"cookie_string": "session=abc", "new_api_user": "7"}, ) class FakeResponse: def __init__(self, payload): self._payload = payload self.cookies = {} self.headers = {"content-type": "application/json"} self.content = b"{}" self.text = "{}" def raise_for_status(self): return None def json(self): return self._payload class FakeHttpClient: def request(self, method, url, **kwargs): path = url.replace("http://newapi.local", "") params = kwargs.get("params") or {} if method == "GET" and path == "/api/token/search": assert params["keyword"] == "SmartUp" return FakeResponse({ "success": True, "data": {"page": 1, "page_size": 100, "total": 0, "items": []}, }) if method == "GET" and path == "/api/token/": return FakeResponse({ "success": True, "data": { "page": 1, "page_size": 100, "total": 4, "items": [ {"id": 447, "name": "SmartUp-4-gptpro-gpt pro", "group": "gpt pro", "key": "sk-XE2o********WWh"}, {"id": 446, "name": "SmartUp-4-gptplus-gpt plus", "group": "gpt plus", "key": "sk-JRi1********rtum"}, {"id": 445, "name": "SmartUp-4-claude特价kiro-claude 特价kiro", "group": "claude 特价kiro", "key": "sk-Aldb********08W2"}, {"id": 56, "name": "plus", "group": "gpt plus", "key": "sk-20cB********pEfE"}, ], }, }) if method == "POST" and path == "/api/token/447/key": return FakeResponse({"success": True, "data": {"key": "sk-gptpro-plain"}}) if method == "POST" and path == "/api/token/446/key": return FakeResponse({"success": True, "data": {"key": "sk-gptplus-plain"}}) if method == "POST" and path == "/api/token/445/key": return FakeResponse({"success": True, "data": {"key": "sk-claude-plain"}}) raise AssertionError(f"unexpected request {method} {path} {params}") client._client = FakeHttpClient() rows = client.list_api_keys(search="SmartUp", status="active") assert [row["id"] for row in rows] == [447, 446, 445] assert [row["group_id"] for row in rows] == ["gpt pro", "gpt plus", "claude 特价kiro"] assert [row["key"] for row in rows] == ["sk-gptpro-plain", "sk-gptplus-plain", "sk-claude-plain"] def test_generated_keys_persists_new_api_tokens_with_plaintext(db_session, monkeypatch): """generated-keys 应把 New-API 远端 token 回填成本地可导入记录。""" from app.routers import upstreams as upstreams_router from app.models.upstream_key import UpstreamGeneratedKey upstream = Upstream( name="NewAPI", base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config_json=json.dumps({"cookie_string": "session=abc", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) class FakeClient: def __init__(self, **kwargs): self.kwargs = kwargs def __enter__(self): return self def __exit__(self, *args): return None def login(self): return None def list_api_keys(self, search="", status="active"): assert search == "SmartUp" return [ { "id": 447, "name": "SmartUp-4-gptpro-gpt pro", "group": "gpt pro", "group_id": "gpt pro", "key": "sk-gptpro-plain", "masked_key": "sk-g********lain", } ] monkeypatch.setattr(upstreams_router, "UpstreamClient", FakeClient) response = upstreams_router.list_generated_keys(upstream.id, db_session, object()) assert len(response) == 1 assert response[0].has_key_value is True assert response[0].id is not None assert response[0].key_name == "SmartUp-4-gptpro-gpt pro" row = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.key_id == "447", ).one() assert row.group_id == "gpt pro" assert row.group_name == "gpt pro" assert row.key_value == "sk-gptpro-plain" assert row.managed_prefix == "SmartUp" def test_generate_keys_allows_new_api_user_upstream(db_session, monkeypatch): """New-API 普通账号上游应允许按分组生成 token。""" from app.routers import upstreams as upstreams_router from app.schemas.upstream import GenerateKeysByGroupsRequest upstream = Upstream( name="NewAPI", base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config_json=json.dumps({"cookie_string": "session=abc", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) monkeypatch.setattr(upstreams_router.website_sync, "reconcile_upstream_keys_full", lambda db, uid: True) class FakeClient: def __init__(self, **kwargs): self.kwargs = kwargs def __enter__(self): return self def __exit__(self, *args): return None def login(self): return None def get_available_groups(self, endpoint): assert endpoint == "/api/user/self/groups" return [{"id": "vip", "name": "VIP"}] def find_smartup_group_key(self, gid, expected_name, prefix): return None def create_api_key(self, name, group_id, **kwargs): assert kwargs["endpoint"] == "/api/token" return {"id": "123", "key": "new-api-plain-key", "masked_key": "new-****-key", "raw": {"id": 123}} monkeypatch.setattr(upstreams_router, "UpstreamClient", FakeClient) response = upstreams_router.generate_keys_by_groups( upstream.id, GenerateKeysByGroupsRequest(group_ids=["vip"], endpoint="/api/token"), db_session, object(), ) assert response.success is True assert response.items[0].status == "created" assert response.items[0].key_value == "new-api-plain-key" # ────────────────────────────────────────────── # Tests for build_new_api_token_name helper # ────────────────────────────────────────────── def test_build_new_api_token_name_fits_50_bytes_with_long_chinese_group(): """长中文分组名 + 长 group_id 生成的 token name 不超过 50 UTF-8 字节。""" from app.services.upstream_client import build_new_api_token_name name = build_new_api_token_name( prefix="SmartUp", upstream_id=7, group_id="claude-max-智能分组-special", group_name="ClaudeMax智能分组SuperLongNameThatWouldBreakOldFormat", ) assert len(name.encode("utf-8")) <= 50, ( f"token name too long: {len(name.encode('utf-8'))} bytes — '{name}'" ) assert name.startswith("SmartUp-7-") def test_build_new_api_token_name_no_broken_multibyte(): """UTF-8 截断后不应出现乱码(decode 不报错)。""" from app.services.upstream_client import build_new_api_token_name name = build_new_api_token_name( prefix="SmartUp", upstream_id=99, group_id="group-x", group_name="中文分组名称超长超长超长超长超长超长超长超长超长超长", ) # 结果必须是有效 UTF-8 字符串,且 ≤ 50 字节 encoded = name.encode("utf-8") assert len(encoded) <= 50 # 没有乱码(可以解回去) assert encoded.decode("utf-8") == name def test_build_new_api_token_name_different_group_ids_produce_different_names(): """两个相同短分组名但不同 group_id 生成不同 token name。""" from app.services.upstream_client import build_new_api_token_name name1 = build_new_api_token_name("SmartUp", 7, "group-a", "智能分组") name2 = build_new_api_token_name("SmartUp", 7, "group-b", "智能分组") assert name1 != name2, "相同分组名但不同 group_id 应生成不同 token name" def test_build_new_api_token_name_prefix_too_long_raises(): """prefix 过长导致骨架超 50 字节时应抛出 ValueError。""" from app.services.upstream_client import build_new_api_token_name # 超长前缀:40 个字符 + upstream_id + separators + hash8 会超 50 very_long_prefix = "A" * 40 with pytest.raises(ValueError, match="过长"): build_new_api_token_name(very_long_prefix, 99, "grp", "name") def test_ensure_group_key_reuses_old_key_name_when_local_record_exists(db_session, monkeypatch): """本地已有旧格式 key_name 记录时,应用旧 key_name 查远端,不重新创建。""" from app.routers.upstreams import _ensure_group_key from app.models.upstream_key import UpstreamGeneratedKey from app.schemas.upstream import GenerateKeysByGroupsRequest # New-API 上游(auth_type=cookie + new_api_user → 使用新短名逻辑) upstream = Upstream( name="NewAPI", base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config_json=json.dumps({"cookie_string": "s=x", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) old_key_name = f"SmartUp-{upstream.id}-ClaudeMax智能分组-claude-max-智能分组" db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="claude-max-智能分组", group_name="ClaudeMax智能分组", key_name=old_key_name, key_value="sk-old", managed_prefix="SmartUp", key_id="remote-old-id", )) db_session.commit() find_calls: list[str] = [] class MockClient: def find_smartup_group_key(self, gid, name, prefix): find_calls.append(name) # 旧格式能找到 if name == old_key_name: return { "id": "remote-old-id", "name": old_key_name, "key": "sk-old-plain", } return None def create_api_key(self, *args, **kwargs): raise AssertionError("create_api_key should NOT be called when old remote key exists") group = {"id": "claude-max-智能分组", "name": "ClaudeMax智能分组"} body = GenerateKeysByGroupsRequest(group_ids=["claude-max-智能分组"], name_prefix="SmartUp", quota=0) result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body) assert result.status == "exists" # 第一次调用用的应该是旧 key_name assert old_key_name in find_calls, f"expected old key_name in find calls, got {find_calls}" def test_ensure_group_key_new_api_uses_short_name_for_create(db_session, monkeypatch): """New-API 上游新建时,发往远端的 token name 是新短名(≤50 字节)。""" from app.routers.upstreams import _ensure_group_key from app.schemas.upstream import GenerateKeysByGroupsRequest from app.services.upstream_client import build_new_api_token_name upstream = Upstream( name="NewAPI", base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config_json=json.dumps({"cookie_string": "s=x", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) gid = "claude-max-智能分组" gname = "ClaudeMax智能分组" expected_name = build_new_api_token_name("SmartUp", upstream.id, gid, gname) created_names: list[str] = [] class MockClient: def find_smartup_group_key(self, group_id, name, prefix): return None # 远端不存在 def create_api_key(self, name, group_id, **kwargs): created_names.append(name) return {"id": "new-123", "key": "sk-new-plain", "masked_key": "sk-****"} group = {"id": gid, "name": gname} body = GenerateKeysByGroupsRequest(group_ids=[gid], name_prefix="SmartUp", quota=0) result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body) assert result.status == "created" assert len(created_names) == 1 assert created_names[0] == expected_name, ( f"expected short name '{expected_name}', got '{created_names[0]}'" ) assert len(expected_name.encode("utf-8")) <= 50 def test_ensure_group_key_prefix_too_long_returns_failed(db_session): """prefix 过长时,_ensure_group_key 应返回 failed 状态,错误信息指向 prefix 过长。""" from app.routers.upstreams import _ensure_group_key from app.schemas.upstream import GenerateKeysByGroupsRequest upstream = Upstream( name="NewAPI", base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config_json=json.dumps({"cookie_string": "s=x", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) very_long_prefix = "MyCompanyLongPrefixName" * 3 # >> 50 bytes class MockClient: def find_smartup_group_key(self, *a, **kw): raise AssertionError("should not reach remote call") def create_api_key(self, *a, **kw): raise AssertionError("should not reach remote call") group = {"id": "vip", "name": "VIP"} body = GenerateKeysByGroupsRequest(group_ids=["vip"], name_prefix=very_long_prefix, quota=0) result = _ensure_group_key(db_session, MockClient(), upstream, group, very_long_prefix, body) assert result.status == "failed" assert "过长" in (result.error or ""), f"expected '过长' in error, got: {result.error}" # ────────────────────────────────────────────── # P1 Regression tests # ────────────────────────────────────────────── def test_p1a_local_key_name_without_key_id_still_queries_remote(db_session): """P1-A: 本地有旧 key_name 但 key_id 为空时,应用旧 key_name 查远端,不跳过。""" from app.routers.upstreams import _ensure_group_key from app.models.upstream_key import UpstreamGeneratedKey from app.schemas.upstream import GenerateKeysByGroupsRequest upstream = Upstream( name="NewAPI", base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config_json=json.dumps({"cookie_string": "s=x", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) old_key_name = f"SmartUp-{upstream.id}-VIP-vip" # 本地记录:有 key_name,但 key_id 为 None db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="vip", group_name="VIP", key_name=old_key_name, key_value="", # 无明文 masked_key="sk-old-masked", managed_prefix="SmartUp", key_id=None, # ← P1-A 的关键:key_id 为空 )) db_session.commit() find_calls: list[str] = [] class MockClient: def find_smartup_group_key(self, gid, name, prefix): find_calls.append(name) if name == old_key_name: return {"id": "remote-777", "name": old_key_name, "key": "sk-plain-found"} return None def create_api_key(self, *args, **kwargs): raise AssertionError("create_api_key should NOT be called when remote key found by old name") group = {"id": "vip", "name": "VIP"} body = GenerateKeysByGroupsRequest(group_ids=["vip"], name_prefix="SmartUp", quota=0) result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body) assert result.status == "exists", f"expected exists, got {result.status}" # 第一次查询必须用旧 key_name assert find_calls[0] == old_key_name, ( f"expected first find call with old_key_name='{old_key_name}', got {find_calls}" ) # 本地记录 key_id 已被回填 row = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.group_id == "vip", ).one() assert row.key_id == "remote-777" assert row.key_value == "sk-plain-found" def test_p1b_no_local_record_old_format_remote_token_imported_not_duplicated(db_session): """P1-B: 无本地记录但远端已有旧格式 token 时,应导入旧 token,不创建新 token。""" from app.routers.upstreams import _ensure_group_key from app.models.upstream_key import UpstreamGeneratedKey from app.schemas.upstream import GenerateKeysByGroupsRequest from app.services.upstream_client import build_new_api_token_name upstream = Upstream( name="NewAPI", base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config_json=json.dumps({"cookie_string": "s=x", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) gid = "vip" gname = "VIP" # 旧格式 token 名(未截断的完整 group_id 拼接) old_format_name = f"SmartUp-{upstream.id}-VIP-{gid}" # 新短名(不同于旧格式) new_short_name = build_new_api_token_name("SmartUp", upstream.id, gid, gname) assert old_format_name != new_short_name, "test setup: old and new names must differ" find_calls: list[str] = [] create_calls: list[str] = [] class MockClient: def find_smartup_group_key(self, group_id, name, prefix): find_calls.append(name) # 仅旧格式能找到 if name == old_format_name: return {"id": "remote-old-888", "name": old_format_name, "key": "sk-old-plain"} return None def create_api_key(self, name, group_id, **kwargs): create_calls.append(name) return {"id": "new-999", "key": "sk-new", "masked_key": "sk-****"} group = {"id": gid, "name": gname} body = GenerateKeysByGroupsRequest(group_ids=[gid], name_prefix="SmartUp", quota=0) result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body) assert result.status == "exists", f"expected exists (imported old token), got {result.status}" assert len(create_calls) == 0, f"create_api_key should not be called, got: {create_calls}" # 本地应只有一条记录,key_name 保留的是旧格式(不覆写为新短名) rows = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.group_id == gid, ).all() assert len(rows) == 1, f"expected exactly 1 local record, got {len(rows)}" assert rows[0].key_id == "remote-old-888" assert rows[0].key_value == "sk-old-plain" assert rows[0].key_name == old_format_name def test_p3_local_old_key_name_synced_to_new_short_name_when_remote_migrated(db_session): """P3: 本地旧 key_name 在远端已消失,远端新短名存在时,本地 key_name 应同步更新为新短名。""" from app.routers.upstreams import _ensure_group_key from app.models.upstream_key import UpstreamGeneratedKey from app.schemas.upstream import GenerateKeysByGroupsRequest from app.services.upstream_client import build_new_api_token_name upstream = Upstream( name="NewAPI", base_url="http://newapi.local", api_prefix="", auth_type="cookie", auth_config_json=json.dumps({"cookie_string": "s=x", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) gid = "vip" gname = "VIP" old_key_name = f"SmartUp-{upstream.id}-VIP-vip" # 旧格式 new_short_name = build_new_api_token_name("SmartUp", upstream.id, gid, gname) assert old_key_name != new_short_name, "test setup: old and new names must differ" # 本地记录用旧名,有 key_id db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id=gid, group_name=gname, key_name=old_key_name, key_value="sk-stale", managed_prefix="SmartUp", key_id="remote-old-id", )) db_session.commit() class MockClient: def find_smartup_group_key(self, group_id, name, prefix): if name == old_key_name: return None # 旧名在远端已消失 if name == new_short_name: return { # 新短名在远端存在 "id": "remote-new-id", "name": new_short_name, "key": "sk-new-plain", } return None def create_api_key(self, *a, **kw): raise AssertionError("create_api_key should not be called") group = {"id": gid, "name": gname} body = GenerateKeysByGroupsRequest(group_ids=[gid], name_prefix="SmartUp", quota=0) result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body) # 远端新短名找到 → status=exists,不创建 assert result.status == "exists", f"expected exists, got {result.status}" # 本地 key_name 应已同步为新短名(P3 一致性) row = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.group_id == gid, ).one() assert row.key_name == new_short_name, ( f"expected key_name updated to '{new_short_name}', got '{row.key_name}'" ) assert row.key_id == "remote-new-id" assert row.key_value == "sk-new-plain" # ────────────────────────────────────────────── # Nox-API Batch / Rate Limit Resilience Tests # ────────────────────────────────────────────── def test_nox_create_token_returns_key_directly(monkeypatch): """Nox POST /api/token/ 直接在 data 中返回 sk-... 时,应直接使用,不查询列表。""" from app.services.upstream_client import UpstreamClient client = UpstreamClient( base_url="http://nox.local", api_prefix="", auth_type="nox_token", auth_config={"token": "abc", "new_api_user": "7"}, ) request_paths = [] def fake_request(method, path, body=None, auth=True): request_paths.append((method, path)) if method == "GET" and path == "/api/status": return {"success": True, "data": {"quota_per_unit": 500000}} if method == "POST" and path == "/api/token/": # Nox-API 直接返回明文 key return {"success": True, "data": {"key": "sk-nox-direct-plaintext-123", "id": 999}} raise AssertionError(f"unexpected request {method} {path}") monkeypatch.setattr(client, "_request", fake_request) # 预热一个空缓存 client._token_list_cache = {} result = client.create_api_key( "SmartUp-7-group-hash", "group-id", quota=2, endpoint="/api/token", ) assert result["id"] == "999" assert result["key"] == "sk-nox-direct-plaintext-123" # 不应该有 GET /api/token/search 或 GET /api/token/ 列表或 POST /api/token/999/key 动作 assert ("GET", "/api/token/") not in request_paths assert ("GET", "/api/token/search") not in request_paths assert ("POST", "/api/token/999/key") not in request_paths def test_nox_create_token_no_key_fallback_success(monkeypatch): """创建响应无 key 时,通过列表和 get-key fallback 补齐。""" from app.services.upstream_client import UpstreamClient client = UpstreamClient( base_url="http://nox.local", api_prefix="", auth_type="nox_token", auth_config={"token": "abc", "new_api_user": "7"}, ) request_paths = [] def fake_request(method, path, body=None, auth=True): request_paths.append((method, path)) if method == "GET" and path == "/api/status": return {"success": True, "data": {"quota_per_unit": 500000}} if method == "POST" and path == "/api/token/": # 不带明文 key return {"success": True, "data": {"id": 888}} if method == "POST" and path == "/api/token/888/key": return {"success": True, "data": {"key": "sk-nox-fallback-key"}} raise AssertionError(f"unexpected request {method} {path}") monkeypatch.setattr(client, "_request", fake_request) # 模拟 _list_new_api_tokens 行为 monkeypatch.setattr( client, "_list_new_api_tokens", lambda search, group_id: [{"id": 888, "name": search, "group": group_id, "key": "sk-****"}] ) result = client.create_api_key( "SmartUp-7-group-hash", "group-id", quota=2, endpoint="/api/token", ) assert result["id"] == "888" assert result["key"] == "sk-nox-fallback-key" def test_nox_create_token_fallback_429_retry(monkeypatch): """fallback 遇 429 时按退避重试成功。""" from app.services.upstream_client import UpstreamClient, UpstreamError import httpx import pytest client = UpstreamClient( base_url="http://nox.local", api_prefix="", auth_type="nox_token", auth_config={"token": "abc", "new_api_user": "7"}, ) request_counts = {"list": 0, "key": 0} def fake_request(method, path, body=None, auth=True): if method == "GET" and path == "/api/status": return {"success": True, "data": {"quota_per_unit": 500000}} if method == "POST" and path == "/api/token/": return {"success": True, "data": {"id": 888}} if method == "POST" and path == "/api/token/888/key": request_counts["key"] += 1 if request_counts["key"] == 1: # 第一次返回 429 response = httpx.Response(429, request=httpx.Request("POST", "http://nox.local")) raise httpx.HTTPStatusError("429 Too Many Requests", request=response.request, response=response) return {"success": True, "data": {"key": "sk-nox-retry-success"}} raise AssertionError(f"unexpected request {method} {path}") monkeypatch.setattr(client, "_request", fake_request) def fake_list(search, group_id): request_counts["list"] += 1 if request_counts["list"] == 1: raise UpstreamError("status code: 429") # 第一次列表调用也模拟 429 return [{"id": 888, "name": search, "group": group_id, "key": "sk-****"}] monkeypatch.setattr(client, "_list_new_api_tokens", fake_list) # 缩短重试延迟以加快测试运行 import app.services.upstream_client monkeypatch.setattr(app.services.upstream_client, "_RATE_LIMIT_BACKOFFS", (0.01, 0.02)) result = client.create_api_key( "SmartUp-7-group-hash", "group-id", quota=2, endpoint="/api/token", ) assert result["id"] == "888" assert result["key"] == "sk-nox-retry-success" assert request_counts["list"] == 2 assert request_counts["key"] == 2 def test_nox_create_token_fallback_fails_raises_pending(monkeypatch): """POST 成功但 fallback 获取 key 429 超限失败时,应抛出 _PendingKeyError。""" from app.services.upstream_client import UpstreamClient, _PendingKeyError import httpx import pytest client = UpstreamClient( base_url="http://nox.local", api_prefix="", auth_type="nox_token", auth_config={"token": "abc", "new_api_user": "7"}, ) def fake_request(method, path, body=None, auth=True): if method == "GET" and path == "/api/status": return {"success": True, "data": {"quota_per_unit": 500000}} if method == "POST" and path == "/api/token/": return {"success": True, "data": {"id": 888}} raise AssertionError(f"unexpected request {method} {path}") monkeypatch.setattr(client, "_request", fake_request) def fake_list(search, group_id): response = httpx.Response(429, request=httpx.Request("GET", "http://nox.local")) raise httpx.HTTPStatusError("429 Too Many Requests", request=response.request, response=response) monkeypatch.setattr(client, "_list_new_api_tokens", fake_list) # 缩短重试延迟以加快测试运行 import app.services.upstream_client monkeypatch.setattr(app.services.upstream_client, "_RATE_LIMIT_BACKOFFS", (0.01, 0.02)) with pytest.raises(_PendingKeyError, match="POST /api/token/ 成功,但查询 token 列表失败"): client.create_api_key( "SmartUp-7-group-hash", "group-id", quota=2, endpoint="/api/token", ) def test_ensure_group_key_nox_batch_uses_cache(db_session, monkeypatch): """验证批量生成 Nox 分组时,使用了缓存,不重复触发 search 请求。""" from app.routers.upstreams import _ensure_group_key from app.schemas.upstream import GenerateKeysByGroupsRequest from app.services.upstream_client import UpstreamClient upstream = Upstream( name="NoxAPI", base_url="http://nox.local", api_prefix="", auth_type="nox_token", auth_config_json=json.dumps({"token": "abc", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) request_calls = [] class MockUpstreamClient(UpstreamClient): def _request(self, method, path, body=None, auth=True): request_calls.append((method, path)) if method == "GET" and path == "/api/status": return {"success": True, "data": {"quota_per_unit": 500000}} if method == "POST" and path == "/api/token/": # 直接在 POST 响应里返回包含 key 的数据以防再次 search/list return {"success": True, "data": {"key": "sk-nox-direct-123", "id": 123}} raise AssertionError(f"unexpected request {method} {path}") client = MockUpstreamClient( base_url="http://nox.local", api_prefix="", auth_type="nox_token", auth_config={"token": "abc", "new_api_user": "7"}, ) # 模拟 pre-warm client._token_list_cache = {} # 批量模拟:2 个不同分组 groups = [ {"id": "grp1", "name": "Group1"}, {"id": "grp2", "name": "Group2"}, ] body = GenerateKeysByGroupsRequest(group_ids=["grp1", "grp2"], name_prefix="SmartUp", quota=0) # 第一次跑 grp1 (应该创建) res1 = _ensure_group_key(db_session, client, upstream, groups[0], "SmartUp", body) assert res1.status == "created" # 第二次跑 grp1 (应该从缓存中匹配 exists,不需要调 search 或任何 HTTP 调用) res2 = _ensure_group_key(db_session, client, upstream, groups[0], "SmartUp", body) assert res2.status == "exists" # 仅在第一次创建时有 POST /api/token/ 请求,获取 status (1次) + 创建 (1次) # 没有针对 exists/exists check 的任何 GET /api/token/search 或 list 调用 get_or_search_calls = [c for c in request_calls if c[0] == "GET" and "token" in c[1]] assert len(get_or_search_calls) == 0, f"Expected no remote list/search calls, got {request_calls}" def test_ensure_group_key_nox_create_fallback_pending_db_save(db_session, monkeypatch): """POST 成功但 fallback 获取 key 429 报错时,本地入库为 created_pending_key,不返回 failed。""" from app.routers.upstreams import _ensure_group_key from app.schemas.upstream import GenerateKeysByGroupsRequest from app.services.upstream_client import _PendingKeyError from app.models.upstream_key import UpstreamGeneratedKey upstream = Upstream( name="NoxAPI", base_url="http://nox.local", api_prefix="", auth_type="nox_token", auth_config_json=json.dumps({"token": "abc", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) class MockClient: def find_smartup_group_key(self, group_id, name, prefix): return None def create_api_key(self, name, group_id, **kwargs): raise _PendingKeyError("simulate list/key 429 failure") group = {"id": "pending-grp", "name": "PendingGrp"} body = GenerateKeysByGroupsRequest(group_ids=["pending-grp"], name_prefix="SmartUp", quota=0) result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body) assert result.status == "created_pending_key" assert "simulate list/key 429 failure" in (result.error or "") # 验证本地数据库保存了记录,但 key_value 为空,以便后续重试 row = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.group_id == "pending-grp", ).one() assert row.status == "created_pending_key" assert row.key_value == "" def test_ensure_group_key_nox_pending_re_invoke_backfills(db_session): """重试 pending 记录时,如果远端已经能查到,应回填 key_id/key_value 并升级为 exists。""" from app.routers.upstreams import _ensure_group_key from app.schemas.upstream import GenerateKeysByGroupsRequest from app.models.upstream_key import UpstreamGeneratedKey upstream = Upstream( name="NoxAPI", base_url="http://nox.local", api_prefix="", auth_type="nox_token", auth_config_json=json.dumps({"token": "abc", "new_api_user": "7"}), groups_endpoint="/api/user/self/groups", rate_endpoint="/api/user/self/groups", ) db_session.add(upstream) db_session.commit() db_session.refresh(upstream) # 模拟本地有一条之前遗留的 pending 记录,key_id 和 key_value 都为空 old_key_name = f"SmartUp-{upstream.id}-PendingGrp-pending-grp" db_session.add(UpstreamGeneratedKey( upstream_id=upstream.id, group_id="pending-grp", group_name="PendingGrp", key_name=old_key_name, key_value="", masked_key="", managed_prefix="SmartUp", key_id=None, status="created_pending_key", )) db_session.commit() class MockClient: def find_smartup_group_key(self, group_id, name, prefix): # 远端此时能够查到了 if name == old_key_name: return {"id": "remote-backfill-id", "name": old_key_name, "key": "sk-backfilled-plain-key"} return None def create_api_key(self, *args, **kwargs): raise AssertionError("create_api_key should not be called for backfilling pending record") group = {"id": "pending-grp", "name": "PendingGrp"} body = GenerateKeysByGroupsRequest(group_ids=["pending-grp"], name_prefix="SmartUp", quota=0) result = _ensure_group_key(db_session, MockClient(), upstream, group, "SmartUp", body) assert result.status == "exists" # 验证本地记录已升级并补全 row = db_session.query(UpstreamGeneratedKey).filter( UpstreamGeneratedKey.upstream_id == upstream.id, UpstreamGeneratedKey.group_id == "pending-grp", ).one() assert row.status == "exists" assert row.key_id == "remote-backfill-id" assert row.key_value == "sk-backfilled-plain-key"