411 lines
15 KiB
Python
411 lines
15 KiB
Python
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.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_original_json_endpoint(db_session, monkeypatch):
|
|
# Setup database rows
|
|
w = Website(
|
|
name="wangwang888",
|
|
site_type="sub2api",
|
|
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" # Numeric string
|
|
)
|
|
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" # Numeric string
|
|
)
|
|
db_session.add(key1)
|
|
db_session.add(key2)
|
|
db_session.commit()
|
|
|
|
updates_recorded = []
|
|
|
|
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",
|
|
"type": "apikey",
|
|
"credentials": {
|
|
"base_url": "old",
|
|
"access_token": "secret_access_token", # Sensitive credential
|
|
"api_key": "old_api_key", # Will be overwritten by local key_value
|
|
"openai_capabilities": "some_capabilities", # Non-sensitive field
|
|
"compact_model_mapping": "some_mapping", # Non-sensitive field
|
|
}
|
|
},
|
|
{"id": "2", "name": "Account 2", "type": "apikey", "credentials": {"base_url": "old"}}
|
|
]
|
|
def extract_id(self, val):
|
|
return val["id"]
|
|
def update_account(self, aid, data):
|
|
updates_recorded.append((aid, data))
|
|
def sync_account_upstream_models(self, aid):
|
|
if aid == "1":
|
|
return ["gpt-4", "gpt-3.5"]
|
|
return [] # 2 returns empty (failure)
|
|
|
|
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)
|
|
# 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
|
|
|
|
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"]
|
|
|
|
assert items_map["2"]["status"] == "failed"
|
|
assert items_map["2"]["model_count"] == 0
|
|
|
|
# Verify update payload regression safety
|
|
# Account 1 should be updated twice (first to fix base_url, second to set model_mapping)
|
|
acc1_updates = [up for up in updates_recorded if up[0] == "1"]
|
|
assert len(acc1_updates) == 2
|
|
|
|
# 1. First update to fix base_url (api_key restored from local key_value="val1")
|
|
payload1 = acc1_updates[0][1]["credentials"]
|
|
assert payload1["base_url"] == "http://up1.api"
|
|
assert "access_token" not in payload1 # Sensitive: still filtered
|
|
assert payload1["api_key"] == "val1" # Restored from db key_value
|
|
assert payload1["openai_capabilities"] == "some_capabilities"
|
|
assert payload1["compact_model_mapping"] == "some_mapping"
|
|
|
|
# 2. Second update to write back model_mapping (api_key still present)
|
|
payload2 = acc1_updates[1][1]["credentials"]
|
|
assert payload2["base_url"] == "http://up1.api"
|
|
assert "access_token" not in payload2
|
|
assert payload2["api_key"] == "val1" # Restored from db key_value
|
|
assert payload2["openai_capabilities"] == "some_capabilities"
|
|
assert payload2["compact_model_mapping"] == "some_mapping"
|
|
assert payload2["model_mapping"] == {"gpt-3.5": "gpt-3.5", "gpt-4": "gpt-4"}
|
|
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
|
|
|
|
def test_sync_upstream_models_streaming_endpoint(db_session, monkeypatch):
|
|
# Setup database rows
|
|
w = Website(
|
|
name="wangwang888",
|
|
site_type="sub2api",
|
|
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()
|
|
|
|
updates_recorded = []
|
|
|
|
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",
|
|
"type": "apikey",
|
|
"credentials": {
|
|
"base_url": "old",
|
|
"access_token": "secret_access_token", # Sensitive credential
|
|
"api_key": "old_api_key", # Will be overwritten by local key_value
|
|
"openai_capabilities": "some_capabilities", # Non-sensitive field
|
|
"compact_model_mapping": "some_mapping", # Non-sensitive field
|
|
}
|
|
},
|
|
{"id": "2", "name": "Account 2", "type": "apikey", "credentials": {"base_url": "old"}}
|
|
]
|
|
def extract_id(self, val):
|
|
return val["id"]
|
|
def update_account(self, aid, data):
|
|
updates_recorded.append((aid, data))
|
|
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
|
|
|
|
# Verify update payload regression safety for stream
|
|
acc1_updates = [up for up in updates_recorded if up[0] == "1"]
|
|
assert len(acc1_updates) == 2
|
|
|
|
# First update to fix base_url (api_key restored from local key_value="val1")
|
|
payload1 = acc1_updates[0][1]["credentials"]
|
|
assert payload1["base_url"] == "http://up1.api"
|
|
assert "access_token" not in payload1 # Sensitive: still filtered
|
|
assert payload1["api_key"] == "val1" # Restored from db key_value
|
|
assert payload1["openai_capabilities"] == "some_capabilities"
|
|
assert payload1["compact_model_mapping"] == "some_mapping"
|
|
|
|
# Second update to write back model_mapping (api_key still present)
|
|
payload2 = acc1_updates[1][1]["credentials"]
|
|
assert payload2["base_url"] == "http://up1.api"
|
|
assert "access_token" not in payload2
|
|
assert payload2["api_key"] == "val1" # Restored from db key_value
|
|
assert payload2["openai_capabilities"] == "some_capabilities"
|
|
assert payload2["compact_model_mapping"] == "some_mapping"
|
|
assert payload2["model_mapping"] == {"gpt-3.5": "gpt-3.5", "gpt-4": "gpt-4"}
|
|
|
|
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()
|
|
|
|
# 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)
|
|
|
|
# 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"]
|
|
|
|
# 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"]
|
|
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
|
|
|
|
def test_sync_upstream_models_empty_key_value_skips_account(db_session, monkeypatch):
|
|
"""本地 key_value 为空时，不允许写回 credentials，跳过该账号并报 failed。"""
|
|
w = Website(
|
|
name="wangwang888",
|
|
site_type="sub2api",
|
|
base_url="https://wangwang.top",
|
|
enabled=True,
|
|
auth_config_json="{}",
|
|
timeout_seconds=30
|
|
)
|
|
db_session.add(w)
|
|
up = Upstream(id=1, name="U1", base_url="http://up1.api", enabled=True)
|
|
db_session.add(up)
|
|
db_session.commit()
|
|
|
|
# key_value 故意留空
|
|
key1 = UpstreamGeneratedKey(
|
|
id=1,
|
|
upstream_id=up.id,
|
|
key_name="k1",
|
|
key_value="", # 空
|
|
group_id="g1",
|
|
group_name="grp1",
|
|
status="active",
|
|
imported_website_id=w.id,
|
|
imported_account_id="1"
|
|
)
|
|
db_session.add(key1)
|
|
db_session.commit()
|
|
|
|
update_calls = []
|
|
|
|
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": "Acc1", "type": "apikey", "credentials": {"base_url": "old", "api_key": "x"}}]
|
|
def extract_id(self, val): return val["id"]
|
|
def update_account(self, aid, data):
|
|
update_calls.append((aid, data))
|
|
def sync_account_upstream_models(self, aid):
|
|
return ["gpt-4"]
|
|
|
|
monkeypatch.setattr("app.routers.websites._client", lambda website: FakeClient())
|
|
|
|
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)
|
|
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
|
|
assert len(data["items"]) == 1
|
|
assert data["items"][0]["status"] == "failed"
|
|
assert "key_value" in data["items"][0]["message"] or "api_key" in data["items"][0]["message"]
|
|
# 最关键：没有调用 update_account（因为在写回之前已提前退出）
|
|
assert len(update_calls) == 0, f"不应调用 update_account，实际调用：{update_calls}"
|
|
finally:
|
|
app.dependency_overrides.clear()
|