271 lines
9.0 KiB
Python
271 lines
9.0 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()
|
|
|
|
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"]
|
|
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
|
|
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()
|
|
|
|
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()
|
|
|
|
# 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()
|