fix: shorten New-API token names
This commit is contained in:
@@ -24,7 +24,7 @@ from app.schemas.upstream import (
|
|||||||
UpstreamCreate, UpstreamUpdate, UpstreamResponse, SnapshotResponse, TestResult,
|
UpstreamCreate, UpstreamUpdate, UpstreamResponse, SnapshotResponse, TestResult,
|
||||||
UpstreamBatchActionItem, UpstreamBatchActionSummary, UpstreamBatchActionResponse,
|
UpstreamBatchActionItem, UpstreamBatchActionSummary, UpstreamBatchActionResponse,
|
||||||
)
|
)
|
||||||
from app.services.upstream_client import UpstreamClient, UpstreamError, build_snapshot, mask_secret, _extract_key_value
|
from app.services.upstream_client import UpstreamClient, UpstreamError, build_snapshot, build_new_api_token_name, mask_secret, _extract_key_value
|
||||||
from app.services.auth_config import MASK, mask_auth_config, normalize_auth_config
|
from app.services.auth_config import MASK, mask_auth_config, normalize_auth_config
|
||||||
from app.services.snapshot_service import diff_snapshots
|
from app.services.snapshot_service import diff_snapshots
|
||||||
from app.services import scheduler as sched_svc
|
from app.services import scheduler as sched_svc
|
||||||
@@ -377,10 +377,27 @@ def _ensure_group_key(
|
|||||||
"""确保一个上游分组有一个 SmartUp 前缀 Key:存在则 upsert,不存在则创建。"""
|
"""确保一个上游分组有一个 SmartUp 前缀 Key:存在则 upsert,不存在则创建。"""
|
||||||
gid = _group_id(group)
|
gid = _group_id(group)
|
||||||
gname = _group_name(group, gid)
|
gname = _group_name(group, gid)
|
||||||
# 使用稳定的 upstream_id + group_id 而非可变名称,避免因改名产生重复
|
|
||||||
# 可读 Key 名:{prefix}-{upstream.id}-{安全的分组名}-{group_id}
|
# 根据上游类型决定 token 名格式:
|
||||||
safe_group_name = re.sub(r"[^a-zA-Z0-9\u4e00-\u9fff_-]", "", gname)[:30] if gname else gid
|
# - New-API/Nox-API:用新短名(hash8,≤50 字节)
|
||||||
stable_name = f"{prefix}-{upstream.id}-{safe_group_name}-{gid}"
|
# - Sub2API 等其他上游:保持旧格式({prefix}-{id}-{safename}-{gid})
|
||||||
|
use_short_name = _is_new_api_user_upstream(upstream) or _is_nox_api_upstream(upstream)
|
||||||
|
if use_short_name:
|
||||||
|
try:
|
||||||
|
stable_name = build_new_api_token_name(prefix, upstream.id, gid, gname)
|
||||||
|
except ValueError as exc:
|
||||||
|
# prefix 过长,骨架已超 50 字节,直接返回清晰错误
|
||||||
|
return GeneratedUpstreamKeyResponse(
|
||||||
|
upstream_id=upstream.id,
|
||||||
|
group_id=gid,
|
||||||
|
group_name=gname,
|
||||||
|
key_name=f"{prefix}-{upstream.id}--??",
|
||||||
|
status="failed",
|
||||||
|
error=str(exc),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
safe_group_name = re.sub(r"[^a-zA-Z0-9\u4e00-\u9fff_-]", "", gname)[:30] if gname else gid
|
||||||
|
stable_name = f"{prefix}-{upstream.id}-{safe_group_name}-{gid}"
|
||||||
|
|
||||||
with _generate_key_lock:
|
with _generate_key_lock:
|
||||||
try:
|
try:
|
||||||
@@ -396,10 +413,18 @@ def _ensure_group_key(
|
|||||||
)
|
)
|
||||||
.first()
|
.first()
|
||||||
)
|
)
|
||||||
if row and row.key_id:
|
if row and row.key_name:
|
||||||
# 本地已有记录,检查远端是否仍存在
|
# 本地已有记录(无论 key_id 是否存在),优先用本地已存的 key_name 查远端
|
||||||
|
# 这样旧格式 key_name 也能找回远端对应 token,不重复创建
|
||||||
|
lookup_name = row.key_name
|
||||||
|
found_by_name = lookup_name # 记录实际命中的远端名称
|
||||||
try:
|
try:
|
||||||
existing = client.find_smartup_group_key(gid, stable_name, prefix)
|
existing = client.find_smartup_group_key(gid, lookup_name, prefix)
|
||||||
|
# 若旧格式未找到,再尝试新格式(同一分组可能已在远端迁移)
|
||||||
|
if existing is None and lookup_name != stable_name:
|
||||||
|
existing = client.find_smartup_group_key(gid, stable_name, prefix)
|
||||||
|
if existing is not None:
|
||||||
|
found_by_name = stable_name
|
||||||
except Exception:
|
except Exception:
|
||||||
existing = None
|
existing = None
|
||||||
if existing:
|
if existing:
|
||||||
@@ -412,6 +437,10 @@ def _ensure_group_key(
|
|||||||
row.masked_key = masked
|
row.masked_key = masked
|
||||||
elif masked:
|
elif masked:
|
||||||
row.masked_key = str(masked)
|
row.masked_key = str(masked)
|
||||||
|
# P3:若是通过新短名找到(旧名已消失),同步更新本地 key_name
|
||||||
|
# 保证 UI 和后续 lookup 与远端名称一致
|
||||||
|
if found_by_name != lookup_name:
|
||||||
|
row.key_name = found_by_name
|
||||||
row.raw_json = json.dumps(existing, ensure_ascii=False)
|
row.raw_json = json.dumps(existing, ensure_ascii=False)
|
||||||
row.status = "exists"
|
row.status = "exists"
|
||||||
row.updated_at = datetime.now(timezone.utc)
|
row.updated_at = datetime.now(timezone.utc)
|
||||||
@@ -422,8 +451,29 @@ def _ensure_group_key(
|
|||||||
# 远端不存在,需要重新创建
|
# 远端不存在,需要重新创建
|
||||||
row.status = "replaced"
|
row.status = "replaced"
|
||||||
|
|
||||||
# 2. 查远端是否有同名 Key(防止并发时另一个请求已创建)
|
# 2. 查远端是否有同名 Key
|
||||||
existing = client.find_smartup_group_key(gid, stable_name, prefix)
|
# 对 New-API/Nox 上游:先按旧格式查(本地无记录但远端可能仍是旧格式),
|
||||||
|
# 再按新短名查(并发创建防重);最后才创建。
|
||||||
|
names_to_search: list[str] = []
|
||||||
|
if use_short_name and not row:
|
||||||
|
# 构造旧格式 token 名,尝试导入已存在的旧格式远端 token
|
||||||
|
old_safe = re.sub(r"[^a-zA-Z0-9\u4e00-\u9fff_-]", "", gname)[:30] if gname else gid
|
||||||
|
old_format_name = f"{prefix}-{upstream.id}-{old_safe}-{gid}"
|
||||||
|
if old_format_name != stable_name:
|
||||||
|
names_to_search.append(old_format_name)
|
||||||
|
names_to_search.append(stable_name)
|
||||||
|
|
||||||
|
existing = None
|
||||||
|
found_name = stable_name
|
||||||
|
for candidate_name in names_to_search:
|
||||||
|
try:
|
||||||
|
existing = client.find_smartup_group_key(gid, candidate_name, prefix)
|
||||||
|
except Exception:
|
||||||
|
existing = None
|
||||||
|
if existing:
|
||||||
|
found_name = candidate_name
|
||||||
|
break
|
||||||
|
|
||||||
if existing:
|
if existing:
|
||||||
key_id = str(existing.get("id") or "")
|
key_id = str(existing.get("id") or "")
|
||||||
key_value = _extract_plaintext_key(existing)
|
key_value = _extract_plaintext_key(existing)
|
||||||
@@ -444,7 +494,7 @@ def _ensure_group_key(
|
|||||||
group_id=gid,
|
group_id=gid,
|
||||||
group_name=gname,
|
group_name=gname,
|
||||||
key_id=key_id or None,
|
key_id=key_id or None,
|
||||||
key_name=stable_name,
|
key_name=found_name,
|
||||||
key_value=key_value or "",
|
key_value=key_value or "",
|
||||||
masked_key=masked,
|
masked_key=masked,
|
||||||
raw_json=json.dumps(existing, ensure_ascii=False),
|
raw_json=json.dumps(existing, ensure_ascii=False),
|
||||||
@@ -456,7 +506,7 @@ def _ensure_group_key(
|
|||||||
db.refresh(row)
|
db.refresh(row)
|
||||||
return _key_response(row, include_value=False)
|
return _key_response(row, include_value=False)
|
||||||
|
|
||||||
# 3. 远端不存在,创建新 Key
|
# 3. 远端不存在,创建新 Key(使用新短名)
|
||||||
created = client.create_api_key(
|
created = client.create_api_key(
|
||||||
stable_name,
|
stable_name,
|
||||||
gid,
|
gid,
|
||||||
@@ -502,8 +552,8 @@ def _ensure_group_key(
|
|||||||
group_name=gname,
|
group_name=gname,
|
||||||
key_name=stable_name,
|
key_name=stable_name,
|
||||||
status="failed",
|
status="failed",
|
||||||
error=str(exc),
|
error=str(exc),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@router.post("/{uid}/keys/generate-by-groups", response_model=GenerateKeysByGroupsResponse)
|
@router.post("/{uid}/keys/generate-by-groups", response_model=GenerateKeysByGroupsResponse)
|
||||||
|
|||||||
@@ -1,7 +1,9 @@
|
|||||||
"""Upstream HTTP client — ported from monitor_ai98pro_group_rates.py."""
|
"""Upstream HTTP client — ported from monitor_ai98pro_group_rates.py."""
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import hashlib
|
||||||
import json
|
import json
|
||||||
|
import re
|
||||||
import time
|
import time
|
||||||
from typing import Any, Callable, Optional
|
from typing import Any, Callable, Optional
|
||||||
from urllib.parse import urljoin
|
from urllib.parse import urljoin
|
||||||
@@ -16,6 +18,60 @@ class UpstreamError(RuntimeError):
|
|||||||
|
|
||||||
NEW_API_DEFAULT_QUOTA_PER_UNIT = 500000
|
NEW_API_DEFAULT_QUOTA_PER_UNIT = 500000
|
||||||
|
|
||||||
|
# New-API/Nox-API token name 上限(UTF-8 字节)
|
||||||
|
_NEW_API_TOKEN_NAME_MAX_BYTES = 50
|
||||||
|
|
||||||
|
|
||||||
|
def _utf8_safe_truncate(text: str, max_bytes: int) -> str:
|
||||||
|
"""按 UTF-8 字节安全截断字符串,不切坏多字节字符。"""
|
||||||
|
encoded = text.encode("utf-8")
|
||||||
|
if len(encoded) <= max_bytes:
|
||||||
|
return text
|
||||||
|
truncated = encoded[:max_bytes]
|
||||||
|
# 向后退直到合法 UTF-8 边界
|
||||||
|
return truncated.decode("utf-8", errors="ignore")
|
||||||
|
|
||||||
|
|
||||||
|
def build_new_api_token_name(
|
||||||
|
prefix: str,
|
||||||
|
upstream_id: int | str,
|
||||||
|
group_id: str,
|
||||||
|
group_name: str,
|
||||||
|
) -> str:
|
||||||
|
"""构造不超过 50 UTF-8 字节的 New-API/Nox-API token 名。
|
||||||
|
|
||||||
|
格式:{prefix}-{upstream_id}-{短分组名}-{hash8}
|
||||||
|
- 短分组名:去除非字母/数字/汉字/_/- 后,UTF-8 安全截断到剩余空间
|
||||||
|
- hash8:sha1("{upstream_id}:{group_id}")[:8],保证同名分组仍能区分
|
||||||
|
- 若 {prefix}-{upstream_id}-{hash8} 本身超过 50 字节,抛出 ValueError
|
||||||
|
|
||||||
|
示例:
|
||||||
|
SmartUp-7-ClaudeMax智能分组-a1b2c3d4
|
||||||
|
"""
|
||||||
|
hash8 = hashlib.sha1(f"{upstream_id}:{group_id}".encode()).hexdigest()[:8]
|
||||||
|
# 最小骨架:{prefix}-{upstream_id}--{hash8}(含两个分隔符)
|
||||||
|
skeleton = f"{prefix}-{upstream_id}--{hash8}"
|
||||||
|
skeleton_bytes = len(skeleton.encode("utf-8"))
|
||||||
|
if skeleton_bytes > _NEW_API_TOKEN_NAME_MAX_BYTES:
|
||||||
|
raise ValueError(
|
||||||
|
f"token name prefix '{prefix}' 过长,骨架 '{skeleton}' 已达 "
|
||||||
|
f"{skeleton_bytes} 字节(上限 {_NEW_API_TOKEN_NAME_MAX_BYTES})"
|
||||||
|
)
|
||||||
|
|
||||||
|
# 可用于短分组名的字节数
|
||||||
|
safe_name = re.sub(r"[^a-zA-Z0-9\u4e00-\u9fff_-]", "", group_name or group_id)
|
||||||
|
# 骨架含一个额外分隔符,留给 group_name 的字节 = max - len(prefix-upid--hash8) - 1
|
||||||
|
# 即 max - skeleton_bytes(skeleton 已包含两个 `-` 和一个 `-` 给 group_name,
|
||||||
|
# 但 skeleton 里 group_name 位置是空的,所以可用字节 = max - skeleton_bytes + 1
|
||||||
|
# skeleton: prefix-uid--hash8 → prefix-uid-<group>-hash8 多出 group 部分
|
||||||
|
# 可用字节:_NEW_API_TOKEN_NAME_MAX_BYTES - len(f"{prefix}-{upstream_id}-".encode()) - len(f"-{hash8}".encode())
|
||||||
|
base_bytes = len(f"{prefix}-{upstream_id}-".encode("utf-8"))
|
||||||
|
suffix_bytes = len(f"-{hash8}".encode("utf-8"))
|
||||||
|
group_budget = _NEW_API_TOKEN_NAME_MAX_BYTES - base_bytes - suffix_bytes
|
||||||
|
short_name = _utf8_safe_truncate(safe_name, group_budget) if group_budget > 0 else ""
|
||||||
|
|
||||||
|
return f"{prefix}-{upstream_id}-{short_name}-{hash8}"
|
||||||
|
|
||||||
|
|
||||||
def _find_token(value: Any) -> str:
|
def _find_token(value: Any) -> str:
|
||||||
if isinstance(value, str) and value.count(".") >= 2:
|
if isinstance(value, str) and value.count(".") >= 2:
|
||||||
|
|||||||
@@ -823,3 +823,399 @@ def test_generate_keys_allows_new_api_user_upstream(db_session, monkeypatch):
|
|||||||
assert response.success is True
|
assert response.success is True
|
||||||
assert response.items[0].status == "created"
|
assert response.items[0].status == "created"
|
||||||
assert response.items[0].key_value == "new-api-plain-key"
|
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"
|
||||||
|
|||||||
Reference in New Issue
Block a user