384 lines
14 KiB
Python
384 lines
14 KiB
Python
"""Grok / XAI 平台识别测试。"""
|
|
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.upstream import Upstream
|
|
from app.models.upstream_key import UpstreamGeneratedKey
|
|
from app.models.snapshot import UpstreamRateSnapshot
|
|
from app.models.website import Website, WebsiteGroupBinding
|
|
from app.routers.websites import (
|
|
_normalize_platform,
|
|
_detect_platform,
|
|
_resolve_platform,
|
|
organize_website_groups,
|
|
import_upstream_keys_as_accounts,
|
|
)
|
|
from app.schemas.website import ImportAccountsRequest
|
|
|
|
|
|
# ── 单元测试：_normalize_platform ─────────────────────
|
|
|
|
def test_normalize_platform_xai_to_grok():
|
|
assert _normalize_platform("xai") == "grok"
|
|
|
|
|
|
def test_normalize_platform_xai_mixed_case():
|
|
assert _normalize_platform("XAI") == "grok"
|
|
|
|
|
|
def test_normalize_platform_grok_preserved():
|
|
assert _normalize_platform("grok") == "grok"
|
|
|
|
|
|
def test_normalize_platform_openai_unchanged():
|
|
assert _normalize_platform("openai") == "openai"
|
|
|
|
|
|
def test_normalize_platform_gemini_unchanged():
|
|
assert _normalize_platform("GEMINI") == "gemini"
|
|
|
|
|
|
def test_normalize_platform_unknown_preserved():
|
|
assert _normalize_platform("some_unknown") == "some_unknown"
|
|
|
|
|
|
# ── 单元测试：_detect_platform ───────────────────────
|
|
|
|
def test_detect_platform_grok_in_name():
|
|
assert _detect_platform("My Grok Key") == "grok"
|
|
|
|
|
|
def test_detect_platform_grok_lowercase():
|
|
assert _detect_platform("grok-group") == "grok"
|
|
|
|
|
|
def test_detect_platform_xai_in_name():
|
|
assert _detect_platform("Chat xAI Key") == "grok"
|
|
|
|
|
|
def test_detect_platform_gemini_still_works():
|
|
assert _detect_platform("Gemini-Pro") == "gemini"
|
|
|
|
|
|
def test_detect_platform_claude_still_works():
|
|
assert _detect_platform("Claude-v2") == "anthropic"
|
|
|
|
|
|
def test_detect_platform_anthropic_still_works():
|
|
assert _detect_platform("anthropic-key") == "anthropic"
|
|
|
|
|
|
def test_detect_platform_antigravity_still_works():
|
|
assert _detect_platform("antigravity-key") == "antigravity"
|
|
|
|
|
|
def test_detect_platform_no_match_fallback():
|
|
assert _detect_platform("abc") == "openai"
|
|
assert _detect_platform("abc", "anthropic") == "anthropic"
|
|
|
|
|
|
def test_detect_platform_fallback_xai_normalized():
|
|
"""fallback='xai' 且未匹配任何关键词 → 规范化返回 grok。"""
|
|
assert _detect_platform("abc", "xai") == "grok"
|
|
|
|
|
|
# ── 单元测试：_resolve_platform ───────────────────────
|
|
|
|
def test_resolve_platform_snapshot_grok():
|
|
"""快照 platform=grok → 返回 grok。"""
|
|
result = _resolve_platform("some key name", {"platform": "grok"}, "openai")
|
|
assert result == "grok"
|
|
|
|
|
|
def test_resolve_platform_snapshot_grok_mixed_case():
|
|
"""快照 platform=Grok → 统一小写返回 grok。"""
|
|
result = _resolve_platform("some key name", {"platform": "Grok"}, "openai")
|
|
assert result == "grok"
|
|
|
|
|
|
def test_resolve_platform_snapshot_xai_normalized():
|
|
"""快照 platform=xai → 规范化为 grok。"""
|
|
result = _resolve_platform("some key name", {"platform": "xai"}, "openai")
|
|
assert result == "grok"
|
|
|
|
|
|
def test_resolve_platform_name_grok_without_snapshot():
|
|
"""无快照平台时，名称中包含 grok → 返回 grok。"""
|
|
result = _resolve_platform("My-Grok-Key", {}, "openai")
|
|
assert result == "grok"
|
|
|
|
|
|
def test_resolve_platform_name_xai_without_snapshot():
|
|
"""无快照平台时，名称中包含 xai → 返回 grok。"""
|
|
result = _resolve_platform("Chat-xAI", {}, "openai")
|
|
assert result == "grok"
|
|
|
|
|
|
def test_resolve_platform_snapshot_overrides_name():
|
|
"""快照有平台时，应优先于名称识别。"""
|
|
result = _resolve_platform("claude-key", {"platform": "grok"}, "openai")
|
|
assert result == "grok", "快照平台应覆盖名称识别"
|
|
|
|
|
|
def test_resolve_platform_gemini_snapshot_still_works():
|
|
"""快照 platform=gemini → 返回 gemini（不受 grok 影响）。"""
|
|
result = _resolve_platform("xxx", {"platform": "gemini"}, "openai")
|
|
assert result == "gemini"
|
|
|
|
|
|
# ── 集成测试夹具 ──────────────────────────────────────
|
|
|
|
@pytest.fixture()
|
|
def db_session():
|
|
engine = create_engine(
|
|
"sqlite://",
|
|
connect_args={"check_same_thread": False},
|
|
poolclass=StaticPool,
|
|
)
|
|
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)
|
|
|
|
|
|
# ── 一键整理：Grok 平台检测与创建 ─────────────────────
|
|
|
|
def test_organize_groups_creates_account_with_grok_platform(db_session, monkeypatch):
|
|
"""一键整理时，快照平台为 grok → 新建账号的 platform = grok。"""
|
|
w = Website(
|
|
name="W1", site_type="sub2api", base_url="http://w1",
|
|
enabled=True, auth_config_json="{}", timeout_seconds=30,
|
|
)
|
|
u1 = Upstream(name="U1", base_url="http://u1")
|
|
db_session.add_all([w, u1])
|
|
db_session.commit()
|
|
db_session.refresh(w)
|
|
db_session.refresh(u1)
|
|
|
|
b1 = WebsiteGroupBinding(
|
|
website_id=w.id, target_group_id="TG1", target_group_name="TG1-Group",
|
|
source_groups_json=json.dumps([{"upstream_id": u1.id, "group_id": "G1"}]),
|
|
enabled=True,
|
|
)
|
|
db_session.add(b1)
|
|
|
|
k1 = UpstreamGeneratedKey(
|
|
upstream_id=u1.id, group_id="G1", group_name="G1-Name",
|
|
key_name="Key-G1", key_value="sk-g1", status="created",
|
|
)
|
|
db_session.add(k1)
|
|
|
|
# 快照明确指定 platform=grok
|
|
snapshot = UpstreamRateSnapshot(
|
|
upstream_id=u1.id,
|
|
snapshot_json=json.dumps({
|
|
"groups": {"G1": {"group_name": "G1-Name", "rate": 0.1, "platform": "grok"}}
|
|
}),
|
|
)
|
|
db_session.add(snapshot)
|
|
db_session.commit()
|
|
|
|
created_bodies = []
|
|
|
|
class MockClient:
|
|
def __init__(self, **kwargs): pass
|
|
def __enter__(self): return self
|
|
def __exit__(self, *a): pass
|
|
def get_groups(self, *a, **kw): return [{"id": "TG1", "name": "TG1-Group"}]
|
|
def list_accounts(self): return []
|
|
def extract_id(self, val): return val.get("id") if isinstance(val, dict) else str(val)
|
|
def create_account(self, body):
|
|
created_bodies.append(body)
|
|
return {"id": "NEW-ACC", "name": body["name"], "group_ids": body["group_ids"]}
|
|
|
|
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
|
|
monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
|
|
|
|
response = organize_website_groups(wid=w.id, db=db_session)
|
|
assert response.success is True
|
|
assert len(created_bodies) == 1
|
|
assert created_bodies[0]["platform"] == "grok"
|
|
|
|
|
|
def test_organize_groups_detects_grok_from_group_name(db_session, monkeypatch):
|
|
"""一键整理时，分组名包含 Grok 且无快照 → 新建账号的 platform = grok。"""
|
|
w = Website(
|
|
name="W1", site_type="sub2api", base_url="http://w1",
|
|
enabled=True, auth_config_json="{}", timeout_seconds=30,
|
|
)
|
|
u1 = Upstream(name="U1", base_url="http://u1")
|
|
db_session.add_all([w, u1])
|
|
db_session.commit()
|
|
db_session.refresh(w)
|
|
db_session.refresh(u1)
|
|
|
|
b1 = WebsiteGroupBinding(
|
|
website_id=w.id, target_group_id="TG1", target_group_name="TG1-Group",
|
|
source_groups_json=json.dumps([{"upstream_id": u1.id, "group_id": "Grok01"}]),
|
|
enabled=True,
|
|
)
|
|
db_session.add(b1)
|
|
|
|
k1 = UpstreamGeneratedKey(
|
|
upstream_id=u1.id, group_id="Grok01", group_name="Grok01-Group",
|
|
key_name="Key-G1", key_value="sk-g1", status="created",
|
|
)
|
|
db_session.add(k1)
|
|
db_session.commit()
|
|
|
|
created_bodies = []
|
|
|
|
class MockClient:
|
|
def __init__(self, **kwargs): pass
|
|
def __enter__(self): return self
|
|
def __exit__(self, *a): pass
|
|
def get_groups(self, *a, **kw): return [{"id": "TG1", "name": "TG1-Group"}]
|
|
def list_accounts(self): return []
|
|
def extract_id(self, val): return val.get("id") if isinstance(val, dict) else str(val)
|
|
def create_account(self, body):
|
|
created_bodies.append(body)
|
|
return {"id": "NEW-ACC", "name": body["name"], "group_ids": body["group_ids"]}
|
|
|
|
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
|
|
monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
|
|
|
|
response = organize_website_groups(wid=w.id, db=db_session)
|
|
assert response.success is True
|
|
assert len(created_bodies) == 1
|
|
assert created_bodies[0]["platform"] == "grok"
|
|
|
|
|
|
def test_organize_groups_corrects_platform_to_grok(db_session, monkeypatch):
|
|
"""已导入账号平台为 openai，快照为 grok → 修正平台为 grok。"""
|
|
w = Website(
|
|
name="W1", site_type="sub2api", base_url="http://w1",
|
|
enabled=True, auth_config_json="{}", timeout_seconds=30,
|
|
)
|
|
u1 = Upstream(name="U1", base_url="http://u1")
|
|
db_session.add_all([w, u1])
|
|
db_session.commit()
|
|
db_session.refresh(w)
|
|
db_session.refresh(u1)
|
|
|
|
b1 = WebsiteGroupBinding(
|
|
website_id=w.id, target_group_id="TG1", target_group_name="TG1-Group",
|
|
source_groups_json=json.dumps([{"upstream_id": u1.id, "group_id": "G1"}]),
|
|
enabled=True,
|
|
)
|
|
db_session.add(b1)
|
|
|
|
k1 = UpstreamGeneratedKey(
|
|
upstream_id=u1.id, group_id="G1", group_name="G1-Group",
|
|
key_name="Key-G1", key_value="sk-g1", status="imported",
|
|
imported_website_id=w.id, imported_account_id="ACC-G1",
|
|
imported_target_group_id="TG1",
|
|
)
|
|
db_session.add(k1)
|
|
|
|
snapshot = UpstreamRateSnapshot(
|
|
upstream_id=u1.id,
|
|
snapshot_json=json.dumps({
|
|
"groups": {"G1": {"group_name": "G1-Group", "rate": 0.1, "platform": "grok"}}
|
|
}),
|
|
)
|
|
db_session.add(snapshot)
|
|
db_session.commit()
|
|
|
|
updated_accounts = []
|
|
mock_platform = "openai"
|
|
|
|
class MockClient:
|
|
def __init__(self, **kwargs): pass
|
|
def __enter__(self): return self
|
|
def __exit__(self, *a): pass
|
|
def get_groups(self, *a, **kw): return [{"id": "TG1", "name": "TG1-Group"}]
|
|
def list_accounts(self):
|
|
return [{"id": "ACC-G1", "name": "SmartUp-G1", "group_ids": ["TG1"], "platform": mock_platform}]
|
|
def extract_id(self, val): return val.get("id") if isinstance(val, dict) else str(val)
|
|
def update_account(self, account_id, body):
|
|
nonlocal mock_platform
|
|
updated_accounts.append((account_id, body))
|
|
if "platform" in body:
|
|
mock_platform = body["platform"]
|
|
return {"id": account_id, "platform": body.get("platform")}
|
|
|
|
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
|
|
monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
|
|
|
|
response = organize_website_groups(wid=w.id, db=db_session)
|
|
assert response.success is True
|
|
assert len(updated_accounts) == 1
|
|
assert updated_accounts[0][1]["platform"] == "grok"
|
|
assert "平台从 openai 修正为 grok" in response.items[0].message
|
|
|
|
|
|
# ── 手动导入：Grok default_platform ───────────────────
|
|
|
|
def test_import_upstream_key_with_grok_platform(monkeypatch, db_session):
|
|
"""手动导入时 default_platform=grok → 创建账号的 platform = grok。"""
|
|
website = Website(
|
|
name="My Sub2API", site_type="sub2api", base_url="http://sub2api.local",
|
|
api_prefix="/api/v1", auth_type="api_key",
|
|
auth_config_json=json.dumps({"key": "admin-key", "header": "x-api-key"}),
|
|
groups_endpoint="/groups", group_update_endpoint="/groups/{id}",
|
|
)
|
|
upstream = Upstream(name="Up1", base_url="http://up1.local")
|
|
db_session.add_all([website, upstream])
|
|
db_session.commit()
|
|
db_session.refresh(website)
|
|
db_session.refresh(upstream)
|
|
|
|
generated = UpstreamGeneratedKey(
|
|
upstream_id=upstream.id, group_id="vip", group_name="VIP",
|
|
key_id="up-key", key_name="SmartUp-VIP", key_value="sk-upstream-generated",
|
|
masked_key="sk-u...", raw_json="{}", status="created",
|
|
)
|
|
db_session.add(generated)
|
|
db_session.commit()
|
|
db_session.refresh(generated)
|
|
|
|
account_bodies = []
|
|
|
|
class FakeClient:
|
|
def __init__(self, **kwargs): pass
|
|
def __enter__(self): return self
|
|
def __exit__(self, *a): pass
|
|
def create_account(self, body, endpoint="/accounts"):
|
|
account_bodies.append(body)
|
|
return {"id": 101, "name": body["name"]}
|
|
def account_exists(self, account_id): return True
|
|
@staticmethod
|
|
def extract_id(data): return str(data.get("id"))
|
|
def get_groups(self, **kw): return []
|
|
def list_accounts(self): return []
|
|
|
|
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", FakeClient)
|
|
monkeypatch.setattr("app.routers.websites.reconcile_upstream_keys_full", lambda db, uid: None)
|
|
monkeypatch.setattr("app.routers.websites.latest_rate_map", lambda db, uid: {})
|
|
monkeypatch.setattr("app.routers.websites.build_target_group_priority_map", lambda db, src: {})
|
|
|
|
response = import_upstream_keys_as_accounts(
|
|
website.id,
|
|
ImportAccountsRequest(
|
|
upstream_key_ids=[generated.id],
|
|
target_group_map={"vip": "7"},
|
|
default_platform="grok",
|
|
platform_mode="manual",
|
|
),
|
|
db_session,
|
|
object(),
|
|
)
|
|
|
|
assert "新建 1" in response.message
|
|
assert len(account_bodies) == 1
|
|
assert account_bodies[0]["platform"] == "grok"
|