1738 lines
69 KiB
Python
1738 lines
69 KiB
Python
"""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"
|
||
|
||
|
||
def test_ensure_group_key_remote_query_failure_aborts(db_session):
|
||
"""验证当远端查询抛出异常时，直接中断，不重建或重置本地记录。"""
|
||
from app.routers.upstreams import _ensure_group_key
|
||
from app.schemas.upstream import GenerateKeysByGroupsRequest
|
||
from app.models.upstream_key import UpstreamGeneratedKey
|
||
from app.services.upstream_client import UpstreamError
|
||
|
||
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)
|
||
|
||
# 1. 本地有记录的情形
|
||
old_key_name = f"SmartUp-{upstream.id}-VIP-vip"
|
||
db_session.add(UpstreamGeneratedKey(
|
||
upstream_id=upstream.id,
|
||
group_id="vip",
|
||
group_name="VIP",
|
||
key_name=old_key_name,
|
||
key_value="sk-local",
|
||
managed_prefix="SmartUp",
|
||
key_id="123",
|
||
status="exists",
|
||
))
|
||
db_session.commit()
|
||
|
||
class MockClient:
|
||
def find_smartup_group_key(self, *args, **kwargs):
|
||
raise UpstreamError("429 Too Many Requests")
|
||
|
||
def create_api_key(self, *args, **kwargs):
|
||
raise AssertionError("create_api_key should not be called on query failure")
|
||
|
||
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 == "failed"
|
||
assert "429 Too Many Requests" in (result.error or "")
|
||
|
||
# 本地状态应该保持 exists, 不应该被置为 replaced
|
||
row = db_session.query(UpstreamGeneratedKey).filter(
|
||
UpstreamGeneratedKey.upstream_id == upstream.id,
|
||
UpstreamGeneratedKey.group_id == "vip",
|
||
).one()
|
||
assert row.status == "exists"
|
||
assert row.key_value == "sk-local"
|
||
|
||
|
||
def test_ensure_group_key_pending_re_invoke_no_key_stays_pending(db_session):
|
||
"""当重试 pending 记录时，如果远端找到 token 但无法获取到明文，应保持 created_pending_key 状态。"""
|
||
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)
|
||
|
||
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):
|
||
# 查到了 token，但无明文 key (例如只返回了 masked_key)
|
||
if name == old_key_name:
|
||
return {"id": "remote-backfill-id", "name": old_key_name, "masked_key": "sk-****"}
|
||
return None
|
||
|
||
def create_api_key(self, *args, **kwargs):
|
||
raise AssertionError("create_api_key should not be called")
|
||
|
||
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 "未能获取到明文 Key" in (result.error or "")
|
||
|
||
# 本地状态应该保持 pending
|
||
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_id == "remote-backfill-id"
|
||
assert row.key_value == ""
|
||
|
||
|
||
def test_generate_keys_by_groups_warm_cache_failure_fails_fast(db_session, monkeypatch):
|
||
"""预热缓存失败时，generate_keys_by_groups 应该直接 abort 抛出 502。"""
|
||
from app.routers.upstreams import generate_keys_by_groups
|
||
from app.schemas.upstream import GenerateKeysByGroupsRequest
|
||
from fastapi import HTTPException
|
||
|
||
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 __init__(self, *args, **kwargs):
|
||
pass
|
||
def __enter__(self):
|
||
return self
|
||
def __exit__(self, *args):
|
||
pass
|
||
def login(self):
|
||
pass
|
||
def get_available_groups(self, endpoint):
|
||
return [{"id": "vip", "name": "VIP"}]
|
||
def warm_token_list_cache(self):
|
||
raise UpstreamError("429 Too Many Requests")
|
||
|
||
from app.routers import upstreams
|
||
monkeypatch.setattr(upstreams, "UpstreamClient", MockClient)
|
||
|
||
body = GenerateKeysByGroupsRequest(group_ids=["vip"], name_prefix="SmartUp", quota=0)
|
||
|
||
with pytest.raises(HTTPException) as excinfo:
|
||
generate_keys_by_groups(upstream.id, body, db=db_session)
|
||
|
||
assert excinfo.value.status_code == 502
|
||
assert "上游连接或缓存预热失败" in excinfo.value.detail
|
||
|