feat(nox-api): implement batch creation rate limit resilience and pending backfilling
This commit is contained in:
liumangmang
@@ -1219,3 +1219,354 @@ def test_p3_local_old_key_name_synced_to_new_short_name_when_remote_migrated(db_
|
||||
)
|
||||
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"
|
||||
|
||||
Reference in New Issue
Block a user