285 lines
10 KiB
Python
285 lines
10 KiB
Python
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.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(
|
||
"sqlite://",
|
||
connect_args={"check_same_thread": False},
|
||
poolclass=StaticPool,
|
||
)
|
||
from app.models import admin_user, upstream, snapshot, webhook_config, notification_log, custom_page, website, revoked_token, upstream_key
|
||
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_sync_upstream_models_workflow(db_session, monkeypatch):
|
||
# 1. 创建网站与上游
|
||
w = Website(
|
||
id=1,
|
||
name="Sub2Api site",
|
||
site_type="sub2api",
|
||
base_url="http://sub2api",
|
||
api_prefix="api/v1",
|
||
auth_type="bearer",
|
||
auth_config_json='{"token": "tok1"}'
|
||
)
|
||
db_session.add(w)
|
||
|
||
up = Upstream(id=10, name="Upstream A", base_url="http://upstream-a")
|
||
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,
|
||
upstream_id=up.id,
|
||
group_id="g1",
|
||
key_name="key-101",
|
||
key_value="val-101",
|
||
status="active",
|
||
imported_website_id=w.id,
|
||
imported_account_id="1001",
|
||
)
|
||
# k2: 属于本站，但账号 ID 无法转为数字
|
||
k2 = UpstreamGeneratedKey(
|
||
id=102,
|
||
upstream_id=up.id,
|
||
group_id="g1",
|
||
key_name="key-102",
|
||
key_value="val-102",
|
||
status="active",
|
||
imported_website_id=w.id,
|
||
imported_account_id="abc",
|
||
)
|
||
# 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.commit()
|
||
|
||
# Mock Sub2Api 客户端
|
||
sync_calls = []
|
||
update_calls = []
|
||
closed_count = 0
|
||
|
||
class MockClient:
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc, tb):
|
||
nonlocal closed_count
|
||
closed_count += 1
|
||
return False
|
||
|
||
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"}},
|
||
]
|
||
|
||
def extract_id(self, val):
|
||
return str(val.get("id"))
|
||
|
||
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 []
|
||
|
||
def update_account(self, account_id, body):
|
||
update_calls.append((account_id, body))
|
||
|
||
monkeypatch.setattr("app.routers.websites._client", lambda row: MockClient())
|
||
|
||
res = sync_website_accounts_upstream_models(wid=w.id, db=db_session)
|
||
|
||
# 验证总体结果
|
||
assert res.success is False # 存在 failed 账号，success 应该为 False
|
||
assert "成功 1 个" in res.message
|
||
assert "失败 3 个" in res.message
|
||
assert "跳过 2 个" in res.message
|
||
|
||
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
|
||
|
||
|
||
def test_sync_upstream_models_list_accounts_none(db_session, monkeypatch):
|
||
w = Website(
|
||
id=1,
|
||
name="Sub2Api site",
|
||
site_type="sub2api",
|
||
base_url="http://sub2api",
|
||
api_prefix="api/v1",
|
||
auth_type="bearer",
|
||
auth_config_json='{"token": "tok1"}'
|
||
)
|
||
db_session.add(w)
|
||
db_session.commit()
|
||
|
||
closed_count = 0
|
||
|
||
class MockClientNone:
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc, tb):
|
||
nonlocal closed_count
|
||
closed_count += 1
|
||
return False
|
||
|
||
def list_accounts(self):
|
||
return None
|
||
|
||
monkeypatch.setattr("app.routers.websites._client", lambda row: MockClientNone())
|
||
|
||
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
|