feat: sync upstream models with real-time streaming updates
This commit is contained in:
SmartUp Developer
@@ -1,15 +1,16 @@
|
||||
import json
|
||||
import pytest
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from sqlalchemy.pool import StaticPool
|
||||
|
||||
from app.database import Base
|
||||
from app.models.website import Website, WebsiteGroupBinding
|
||||
from app.models.upstream import Upstream
|
||||
from app.models.upstream_key import UpstreamGeneratedKey
|
||||
from app.models.website import Website
|
||||
from app.routers.websites import sync_website_accounts_upstream_models
|
||||
from app.schemas.website import SyncUpstreamModelsResponse
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def db_session():
|
||||
engine = create_engine(
|
||||
@@ -28,257 +29,242 @@ def db_session():
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
|
||||
|
||||
def test_sync_upstream_models_workflow(db_session, monkeypatch):
|
||||
# 1. 创建网站与上游
|
||||
def test_sync_upstream_models_original_json_endpoint(db_session, monkeypatch):
|
||||
# Setup database rows
|
||||
w = Website(
|
||||
id=1,
|
||||
name="Sub2Api site",
|
||||
name="wangwang888",
|
||||
site_type="sub2api",
|
||||
base_url="http://sub2api",
|
||||
api_prefix="api/v1",
|
||||
auth_type="bearer",
|
||||
auth_config_json='{"token": "tok1"}'
|
||||
base_url="https://wangwang.top",
|
||||
enabled=True,
|
||||
auth_config_json="{}",
|
||||
timeout_seconds=30
|
||||
)
|
||||
db_session.add(w)
|
||||
|
||||
up = Upstream(id=10, name="Upstream A", base_url="http://upstream-a")
|
||||
up = Upstream(id=1, name="Upstream 1", base_url="http://up1.api", enabled=True)
|
||||
db_session.add(up)
|
||||
|
||||
up_empty = Upstream(id=11, name="Upstream Empty", base_url="")
|
||||
db_session.add(up_empty)
|
||||
db_session.commit()
|
||||
|
||||
# 2. 插入本地导入的 Key 记录
|
||||
# k1: 属于本站，且在远端正常存在的账号
|
||||
k1 = UpstreamGeneratedKey(
|
||||
id=101,
|
||||
key1 = UpstreamGeneratedKey(
|
||||
id=1,
|
||||
upstream_id=up.id,
|
||||
key_name="k1",
|
||||
key_value="val1",
|
||||
group_id="g1",
|
||||
key_name="key-101",
|
||||
key_value="val-101",
|
||||
group_name="grp1",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1001",
|
||||
imported_account_id="1" # Numeric string
|
||||
)
|
||||
# k2: 属于本站，但账号 ID 无法转为数字
|
||||
k2 = UpstreamGeneratedKey(
|
||||
id=102,
|
||||
key2 = UpstreamGeneratedKey(
|
||||
id=2,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-102",
|
||||
key_value="val-102",
|
||||
key_name="k2",
|
||||
key_value="val2",
|
||||
group_id="g2",
|
||||
group_name="grp2",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="abc",
|
||||
imported_account_id="2" # Numeric string
|
||||
)
|
||||
# k3: 属于其他站的账号
|
||||
k3 = UpstreamGeneratedKey(
|
||||
id=103,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-103",
|
||||
key_value="val-103",
|
||||
status="active",
|
||||
imported_website_id=999,
|
||||
imported_account_id="1003",
|
||||
)
|
||||
# k4: 属于本站，但远端已被删除的账号
|
||||
k4 = UpstreamGeneratedKey(
|
||||
id=104,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-104",
|
||||
key_value="val-104",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1004",
|
||||
)
|
||||
# k5: 属于本站，但同步接口返回空模型（应当标记为 failed，但应先安全修复 base_url）
|
||||
k5 = UpstreamGeneratedKey(
|
||||
id=105,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-105",
|
||||
key_value="val-105",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1005",
|
||||
)
|
||||
# k6: 属于本站，但同步接口抛错（应当标记为 failed，但应先安全修复 base_url）
|
||||
k6 = UpstreamGeneratedKey(
|
||||
id=106,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-106",
|
||||
key_value="val-106",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1006",
|
||||
)
|
||||
# k7: 属于本站，但对应上游 base_url 为空（应当直接标记为 failed 且不调用同步和写回）
|
||||
k7 = UpstreamGeneratedKey(
|
||||
id=107,
|
||||
upstream_id=up_empty.id,
|
||||
group_id="g1",
|
||||
key_name="key-107",
|
||||
key_value="val-107",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1007",
|
||||
)
|
||||
db_session.add_all([k1, k2, k3, k4, k5, k6, k7])
|
||||
db_session.add(key1)
|
||||
db_session.add(key2)
|
||||
db_session.commit()
|
||||
|
||||
# Mock Sub2Api 客户端
|
||||
sync_calls = []
|
||||
update_calls = []
|
||||
closed_count = 0
|
||||
|
||||
class MockClient:
|
||||
class FakeClient:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
nonlocal closed_count
|
||||
closed_count += 1
|
||||
return False
|
||||
|
||||
def __exit__(self, *args):
|
||||
pass
|
||||
def list_accounts(self):
|
||||
return [
|
||||
{"id": 1001, "name": "acc-1001", "credentials": {"api_key": "k1", "model_mapping": {"old": "old"}, "base_url": "http://default-base-url", "compact_model_mapping": {"compact": "compact"}, "openai_capabilities": "caps"}},
|
||||
{"id": "abc", "name": "acc-abc", "credentials": {}},
|
||||
{"id": 1005, "name": "acc-1005", "credentials": {"api_key": "k5"}},
|
||||
{"id": 1006, "name": "acc-1006", "credentials": {"api_key": "k6"}},
|
||||
{"id": 1007, "name": "acc-1007", "credentials": {"api_key": "k7"}},
|
||||
{"id": "1", "name": "Account 1", "credentials": {"base_url": "old"}},
|
||||
{"id": "2", "name": "Account 2", "credentials": {"base_url": "old"}}
|
||||
]
|
||||
|
||||
def extract_id(self, val):
|
||||
return str(val.get("id"))
|
||||
return val["id"]
|
||||
def update_account(self, aid, data):
|
||||
pass
|
||||
def sync_account_upstream_models(self, aid):
|
||||
if aid == "1":
|
||||
return ["gpt-4", "gpt-3.5"]
|
||||
return [] # 2 returns empty (failure)
|
||||
|
||||
def sync_account_upstream_models(self, account_id):
|
||||
sync_calls.append(account_id)
|
||||
if account_id == "1001":
|
||||
return ["gpt-4", "gpt-3.5-turbo", "", "gpt-4 "]
|
||||
elif account_id == "1005":
|
||||
return []
|
||||
elif account_id == "1006":
|
||||
raise Exception("Network Error")
|
||||
return []
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda website: FakeClient())
|
||||
|
||||
def update_account(self, account_id, body):
|
||||
update_calls.append((account_id, body))
|
||||
# Override dependencies
|
||||
from fastapi.testclient import TestClient
|
||||
from app.main import app
|
||||
from app.database import get_db
|
||||
from app.utils.auth import get_current_user
|
||||
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda row: MockClient())
|
||||
app.dependency_overrides[get_db] = lambda: db_session
|
||||
app.dependency_overrides[get_current_user] = lambda: None
|
||||
|
||||
res = sync_website_accounts_upstream_models(wid=w.id, db=db_session)
|
||||
try:
|
||||
client = TestClient(app)
|
||||
# 1. Verify original JSON POST endpoint behaves exactly as before
|
||||
resp = client.post(f"/api/websites/{w.id}/accounts/sync-upstream-models")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["success"] is False # Because one account failed (2 returned empty models list)
|
||||
assert len(data["items"]) == 2
|
||||
|
||||
# 验证总体结果
|
||||
assert res.success is False # 存在 failed 账号，success 应该为 False
|
||||
assert "成功 1 个" in res.message
|
||||
assert "失败 3 个" in res.message
|
||||
assert "跳过 2 个" in res.message
|
||||
items_map = {item["account_id"]: item for item in data["items"]}
|
||||
assert items_map["1"]["status"] == "success"
|
||||
assert items_map["1"]["model_count"] == 2
|
||||
assert items_map["1"]["models"] == ["gpt-3.5", "gpt-4"]
|
||||
|
||||
items = res.items
|
||||
assert len(items) == 6
|
||||
|
||||
# 1001 成功 (修复 Base URL 并同步)
|
||||
item_1001 = next(i for i in items if i.account_id == "1001")
|
||||
assert item_1001.status == "success"
|
||||
assert item_1001.model_count == 2
|
||||
assert item_1001.models == ["gpt-3.5-turbo", "gpt-4"]
|
||||
assert "已修复 Base URL" in item_1001.message
|
||||
|
||||
# 1001 的两次 update 调用检测
|
||||
calls_1001 = [call for call in update_calls if call[0] == "1001"]
|
||||
assert len(calls_1001) == 2
|
||||
# 第一次安全修复：应当在保留其他非敏感 credentials (如 compact_model_mapping) 的基础上，只将 base_url 变更为 http://upstream-a，并且过滤掉敏感字段 api_key
|
||||
assert calls_1001[0][1]["credentials"]["base_url"] == "http://upstream-a"
|
||||
assert calls_1001[0][1]["credentials"]["compact_model_mapping"] == {"compact": "compact"}
|
||||
assert calls_1001[0][1]["credentials"]["openai_capabilities"] == "caps"
|
||||
assert "api_key" not in calls_1001[0][1]["credentials"]
|
||||
# 第二次同步成功写入：同样保留其他非敏感 credentials，过滤掉敏感字段 api_key，将 model_mapping 更新进去
|
||||
assert calls_1001[1][1]["credentials"]["base_url"] == "http://upstream-a"
|
||||
assert calls_1001[1][1]["credentials"]["compact_model_mapping"] == {"compact": "compact"}
|
||||
assert calls_1001[1][1]["credentials"]["openai_capabilities"] == "caps"
|
||||
assert "api_key" not in calls_1001[1][1]["credentials"]
|
||||
assert calls_1001[1][1]["credentials"]["model_mapping"] == {
|
||||
"gpt-3.5-turbo": "gpt-3.5-turbo",
|
||||
"gpt-4": "gpt-4"
|
||||
}
|
||||
|
||||
# abc 跳过
|
||||
item_abc = next(i for i in items if i.account_id == "abc")
|
||||
assert item_abc.status == "skipped"
|
||||
assert "非数字" in item_abc.message
|
||||
|
||||
# 1004 跳过
|
||||
item_1004 = next(i for i in items if i.account_id == "1004")
|
||||
assert item_1004.status == "skipped"
|
||||
assert "不存在" in item_1004.message
|
||||
|
||||
# 1005 失败 (空模型，但也应当执行了第一次安全修复的 update_account)
|
||||
item_1005 = next(i for i in items if i.account_id == "1005")
|
||||
assert item_1005.status == "failed"
|
||||
assert "已修复 Base URL" in item_1005.message
|
||||
assert "为空" in item_1005.message
|
||||
calls_1005 = [call for call in update_calls if call[0] == "1005"]
|
||||
assert len(calls_1005) == 1
|
||||
assert calls_1005[0][1]["credentials"]["base_url"] == "http://upstream-a"
|
||||
assert "api_key" not in calls_1005[0][1]["credentials"]
|
||||
|
||||
# 1006 失败 (抛错，但应当执行了第一次安全修复的 update_account)
|
||||
item_1006 = next(i for i in items if i.account_id == "1006")
|
||||
assert item_1006.status == "failed"
|
||||
assert "已修复 Base URL" in item_1006.message
|
||||
assert "Network Error" in item_1006.message
|
||||
calls_1006 = [call for call in update_calls if call[0] == "1006"]
|
||||
assert len(calls_1006) == 1
|
||||
assert calls_1006[0][1]["credentials"]["base_url"] == "http://upstream-a"
|
||||
assert "api_key" not in calls_1006[0][1]["credentials"]
|
||||
|
||||
# 1007 失败 (上游 base_url 为空，不应当调用 update_account 且不调用同步模型)
|
||||
item_1007 = next(i for i in items if i.account_id == "1007")
|
||||
assert item_1007.status == "failed"
|
||||
assert "base_url 为空" in item_1007.message
|
||||
calls_1007 = [call for call in update_calls if call[0] == "1007"]
|
||||
assert len(calls_1007) == 0
|
||||
assert "1007" not in sync_calls
|
||||
|
||||
assert closed_count == 1
|
||||
assert items_map["2"]["status"] == "failed"
|
||||
assert items_map["2"]["model_count"] == 0
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
def test_sync_upstream_models_list_accounts_none(db_session, monkeypatch):
|
||||
def test_sync_upstream_models_streaming_endpoint(db_session, monkeypatch):
|
||||
# Setup database rows
|
||||
w = Website(
|
||||
id=1,
|
||||
name="Sub2Api site",
|
||||
name="wangwang888",
|
||||
site_type="sub2api",
|
||||
base_url="http://sub2api",
|
||||
api_prefix="api/v1",
|
||||
auth_type="bearer",
|
||||
auth_config_json='{"token": "tok1"}'
|
||||
base_url="https://wangwang.top",
|
||||
enabled=True,
|
||||
auth_config_json="{}",
|
||||
timeout_seconds=30
|
||||
)
|
||||
db_session.add(w)
|
||||
up = Upstream(id=1, name="Upstream 1", base_url="http://up1.api", enabled=True)
|
||||
db_session.add(up)
|
||||
db_session.commit()
|
||||
|
||||
key1 = UpstreamGeneratedKey(
|
||||
id=1,
|
||||
upstream_id=up.id,
|
||||
key_name="k1",
|
||||
key_value="val1",
|
||||
group_id="g1",
|
||||
group_name="grp1",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1"
|
||||
)
|
||||
key2 = UpstreamGeneratedKey(
|
||||
id=2,
|
||||
upstream_id=up.id,
|
||||
key_name="k2",
|
||||
key_value="val2",
|
||||
group_id="g2",
|
||||
group_name="grp2",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="2"
|
||||
)
|
||||
db_session.add(key1)
|
||||
db_session.add(key2)
|
||||
db_session.commit()
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
def __enter__(self):
|
||||
return self
|
||||
def __exit__(self, *args):
|
||||
pass
|
||||
def list_accounts(self):
|
||||
return [
|
||||
{"id": "1", "name": "Account 1", "credentials": {"base_url": "old"}},
|
||||
{"id": "2", "name": "Account 2", "credentials": {"base_url": "old"}}
|
||||
]
|
||||
def extract_id(self, val):
|
||||
return val["id"]
|
||||
def update_account(self, aid, data):
|
||||
pass
|
||||
def sync_account_upstream_models(self, aid):
|
||||
if aid == "1":
|
||||
return ["gpt-4", "gpt-3.5"]
|
||||
raise Exception("connection timed out") # Connection failure for 2
|
||||
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda website: FakeClient())
|
||||
|
||||
# Override dependencies
|
||||
from fastapi.testclient import TestClient
|
||||
from app.main import app
|
||||
from app.database import get_db
|
||||
from app.utils.auth import get_current_user
|
||||
|
||||
app.dependency_overrides[get_db] = lambda: db_session
|
||||
app.dependency_overrides[get_current_user] = lambda: None
|
||||
|
||||
try:
|
||||
client = TestClient(app)
|
||||
# 2. Verify new Streaming endpoint returning ndjson lines
|
||||
resp = client.post(f"/api/websites/{w.id}/accounts/sync-upstream-models/stream")
|
||||
assert resp.status_code == 200
|
||||
assert "application/x-ndjson" in resp.headers["content-type"]
|
||||
|
||||
lines = [line if isinstance(line, str) else line.decode("utf-8") for line in resp.iter_lines() if line]
|
||||
parsed_events = [json.loads(line) for line in lines]
|
||||
|
||||
assert len(parsed_events) == 4
|
||||
# Verify event sequence
|
||||
assert parsed_events[0]["event"] == "start"
|
||||
assert parsed_events[0]["data"]["total_accounts"] == 2
|
||||
|
||||
assert parsed_events[1]["event"] == "item"
|
||||
assert parsed_events[1]["data"]["account_id"] == "1"
|
||||
assert parsed_events[1]["data"]["status"] == "success"
|
||||
assert parsed_events[1]["data"]["model_count"] == 2
|
||||
|
||||
assert parsed_events[2]["event"] == "item"
|
||||
assert parsed_events[2]["data"]["account_id"] == "2"
|
||||
assert parsed_events[2]["data"]["status"] == "failed"
|
||||
assert "connection timed out" in parsed_events[2]["data"]["message"]
|
||||
|
||||
assert parsed_events[3]["event"] == "complete"
|
||||
assert parsed_events[3]["data"]["success"] is False
|
||||
assert len(parsed_events[3]["data"]["items"]) == 2
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
def test_sync_upstream_models_error_handling(db_session, monkeypatch):
|
||||
# Setup website with wrong site_type
|
||||
w = Website(
|
||||
name="W1",
|
||||
site_type="invalid_type", # Should cause validation error
|
||||
base_url="https://wangwang.top",
|
||||
enabled=True,
|
||||
auth_config_json="{}",
|
||||
timeout_seconds=30
|
||||
)
|
||||
db_session.add(w)
|
||||
db_session.commit()
|
||||
|
||||
closed_count = 0
|
||||
# Override dependencies
|
||||
from fastapi.testclient import TestClient
|
||||
from app.main import app
|
||||
from app.database import get_db
|
||||
from app.utils.auth import get_current_user
|
||||
|
||||
class MockClientNone:
|
||||
def __enter__(self):
|
||||
return self
|
||||
app.dependency_overrides[get_db] = lambda: db_session
|
||||
app.dependency_overrides[get_current_user] = lambda: None
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
nonlocal closed_count
|
||||
closed_count += 1
|
||||
return False
|
||||
try:
|
||||
client = TestClient(app)
|
||||
|
||||
def list_accounts(self):
|
||||
return None
|
||||
# Non-streaming endpoint returns error response for invalid site type
|
||||
resp = client.post(f"/api/websites/{w.id}/accounts/sync-upstream-models")
|
||||
assert resp.status_code == 400
|
||||
assert "only sub2api" in resp.json()["detail"]
|
||||
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda row: MockClientNone())
|
||||
# Streaming endpoint yields "error" event and exits
|
||||
resp_stream = client.post(f"/api/websites/{w.id}/accounts/sync-upstream-models/stream")
|
||||
assert resp_stream.status_code == 200
|
||||
lines = [line if isinstance(line, str) else line.decode("utf-8") for line in resp_stream.iter_lines() if line]
|
||||
parsed_events = [json.loads(line) for line in lines]
|
||||
assert len(parsed_events) == 1
|
||||
assert parsed_events[0]["event"] == "error"
|
||||
assert "only sub2api" in parsed_events[0]["data"]["message"]
|
||||
|
||||
res = sync_website_accounts_upstream_models(wid=w.id, db=db_session)
|
||||
|
||||
assert res.success is False
|
||||
assert "拉取远端账号列表失败" in res.message
|
||||
assert len(res.items) == 0
|
||||
assert closed_count == 1
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
Reference in New Issue
Block a user