fix: reuse upstream keys for account import

This commit is contained in:
SmartUp Developer
2026-05-24 23:18:40 +08:00
parent 6044b00685
commit 3a31d185a4
6 changed files with 253 additions and 27 deletions
+123
View File
@@ -262,6 +262,64 @@ def test_ensure_group_key_reuses_old_record(db_session, monkeypatch):
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
@@ -310,6 +368,71 @@ def test_sync_removes_remote_key_when_list_empty(db_session, monkeypatch):
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