"""E2E test for Endpoint Format Restriction feature. 完整测试 6 个阶段: 0. 环境发现(自动查找渠道和模型) 1. 功能关闭 - 所有跨格式请求不被拦截 2. 功能开启 - 匹配的渠道+端点放行 3. 功能开启 - 不匹配的渠道+端点被拦截 4. 错误消息验证 5. 清理 """ import argparse import json import sys import time import mysql.connector import requests DEFAULT_BASE_URL = "http://localhost:3000" DEFAULT_DB_HOST = "localhost" DEFAULT_DB_PORT = 33306 DEFAULT_DB_USER = "root" DEFAULT_DB_PASS = "Lanqi123456" DEFAULT_DB_NAME = "new-api" OPTION_KEY = "global.endpoint_format_restriction_enabled" # 渠道类型常量 (与 Go constant 对应) CH_OPENAI = 1 CH_MIDJOURNEY = 2 CH_AZURE = 3 CH_OLLAMA = 4 CH_ANTHROPIC = 14 CH_GEMINI = 24 CH_DEEPSEEK = 43 CH_VERTEX_AI = 41 CH_ALI = 17 CH_AWS = 33 CH_KLING = 50 CH_SUNO = 36 # 渠道类型 → 默认格式家族 CHANNEL_FAMILIES = { CH_OPENAI: ["openai"], CH_AZURE: ["openai"], CH_DEEPSEEK: ["openai"], CH_ANTHROPIC: ["claude"], CH_AWS: ["claude"], CH_GEMINI: ["gemini"], CH_VERTEX_AI: ["openai", "claude", "gemini"], CH_ALI: ["openai", "claude"], } # 非LLM渠道(不在 channelNativeFamilies 中) NON_LLM_CHANNELS = [CH_MIDJOURNEY, CH_KLING, CH_SUNO] def get_db(args): return mysql.connector.connect( host=args.db_host, port=args.db_port, user=args.db_user, password=args.db_pass, database=args.db_name, ) def get_token_key(args) -> str: db = get_db(args) cur = db.cursor() cur.execute("SELECT `key` FROM tokens WHERE id = %s", (args.token_id,)) row = cur.fetchone() cur.close() db.close() if not row: raise ValueError(f"Token id={args.token_id} not found") return row[0] def set_restriction(args, enabled: bool): db = get_db(args) cur = db.cursor() val = "true" if enabled else "false" cur.execute( "INSERT INTO options (`key`, value) VALUES (%s, %s) ON DUPLICATE KEY UPDATE value = %s", (OPTION_KEY, val, val), ) db.commit() cur.close() db.close() def discover_channels(args) -> dict: """从 DB 查询可用渠道,按 type 分组。返回 {type: [(id, model_name), ...]}""" db = get_db(args) cur = db.cursor() cur.execute("SELECT id, type, models FROM channels WHERE status = 1") rows = cur.fetchall() cur.close() db.close() result = {} for ch_id, ch_type, models_str in rows: if not models_str: continue models = [m.strip() for m in models_str.split(",") if m.strip()] if models: result.setdefault(ch_type, []).append((ch_id, models[0])) return result def find_channel(channels: dict, ch_type: int) -> tuple: """查找指定类型的渠道,返回 (channel_id, model_name) 或 None""" entries = channels.get(ch_type) if entries: return entries[0] return None # ============================================================================ # Request helpers # ============================================================================ def chat_req(key: str, model: str) -> requests.Response: return requests.post( f"{BASE_URL}/v1/chat/completions", headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"}, json={"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5}, timeout=30, ) def messages_req(key: str, model: str) -> requests.Response: return requests.post( f"{BASE_URL}/v1/messages", headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"}, json={"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5}, timeout=30, ) def gemini_req(key: str, model: str) -> requests.Response: return requests.post( f"{BASE_URL}/v1beta/models/{model}:generateContent", headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"}, json={"contents": [{"parts": [{"text": "hi"}]}]}, timeout=30, ) def is_blocked(resp: requests.Response) -> bool: if resp.status_code != 400: return False try: return resp.json().get("error", {}).get("code") == "endpoint_format_restricted" except Exception: return False def wait_for_restriction(args, token_key: str, model: str, enabled: bool) -> bool: """轮询等待限制设置生效(最长 70s,每 3s 检查一次)""" for i in range(24): r = chat_req(token_key, model) blocked = is_blocked(r) if enabled and blocked: print(f" Restriction ACTIVE (after ~{i*3}s)") return True if not enabled and not blocked: print(f" Restriction INACTIVE (after ~{i*3}s)") return True time.sleep(3) print(f" WARNING: Restriction state not confirmed after 70s") return False # ============================================================================ # Test runners # ============================================================================ BASE_URL = DEFAULT_BASE_URL passed = 0 failed = 0 def run_test(name: str, fn) -> bool: global passed, failed try: fn() print(f" [PASS] {name}") passed += 1 return True except AssertionError as e: print(f" [FAIL] {name}: {e}") failed += 1 return False def phase1_off(token_key, openai_model, claude_model, gemini_info): """阶段 1: 功能关闭 - 所有跨格式请求不被拦截""" print("\n--- 阶段 1: 功能关闭 ---") def test_cross_openai_claude(): r = chat_req(token_key, claude_model) assert not is_blocked(r), f"Claude模型 + OpenAI端点 不应被拦截, got {r.status_code}: {r.text[:200]}" def test_cross_claude_openai(): r = messages_req(token_key, openai_model) assert not is_blocked(r), f"OpenAI模型 + Claude端点 不应被拦截, got {r.status_code}: {r.text[:200]}" def test_cross_gemini_openai(): if gemini_info: r = chat_req(token_key, gemini_info[1]) assert not is_blocked(r), f"Gemini模型 + OpenAI端点 不应被拦截" run_test("Claude模型 + OpenAI端点 不拦截", test_cross_openai_claude) run_test("OpenAI模型 + Claude端点 不拦截", test_cross_claude_openai) run_test("Gemini模型 + OpenAI端点 不拦截", test_cross_gemini_openai) def phase2_allowed(token_key, openai_model, claude_model, gemini_info, deepseek_info): """阶段 2: 功能开启 - 匹配的渠道+端点放行""" print("\n--- 阶段 2: 功能开启 - 匹配放行 ---") def test_openai_openai(): r = chat_req(token_key, openai_model) assert not is_blocked(r), f"OpenAI渠道 + OpenAI端点 应放行, got {r.status_code}" def test_claude_claude(): r = messages_req(token_key, claude_model) assert not is_blocked(r), f"Claude渠道 + Claude端点 应放行, got {r.status_code}" def test_gemini_gemini(): if gemini_info: r = gemini_req(token_key, gemini_info[1]) assert not is_blocked(r), f"Gemini渠道 + Gemini端点 应放行" def test_deepseek_openai(): if deepseek_info: r = chat_req(token_key, deepseek_info[1]) assert not is_blocked(r), f"DeepSeek渠道 + OpenAI端点 应放行" run_test("OpenAI渠道 + /v1/chat/completions 放行", test_openai_openai) run_test("Claude渠道 + /v1/messages 放行", test_claude_claude) run_test("Gemini渠道 + /v1beta 放行", test_gemini_gemini) run_test("DeepSeek渠道 + /v1/chat/completions 放行", test_deepseek_openai) def phase3_blocked(token_key, openai_model, claude_model, gemini_info, deepseek_info): """阶段 3: 功能开启 - 不匹配的渠道+端点被拦截""" print("\n--- 阶段 3: 功能开启 - 不匹配拦截 ---") def test_claude_openai_blocked(): r = chat_req(token_key, claude_model) assert is_blocked(r), f"Claude渠道 + OpenAI端点 应被拦截, got {r.status_code}: {r.text[:200]}" def test_openai_claude_blocked(): r = messages_req(token_key, openai_model) assert is_blocked(r), f"OpenAI渠道 + Claude端点 应被拦截, got {r.status_code}: {r.text[:200]}" def test_gemini_openai_blocked(): if gemini_info: r = chat_req(token_key, gemini_info[1]) assert is_blocked(r), f"Gemini渠道 + OpenAI端点 应被拦截" def test_openai_gemini_blocked(): if gemini_info: r = gemini_req(token_key, openai_model) assert is_blocked(r), f"OpenAI渠道 + Gemini端点 应被拦截" def test_deepseek_claude_blocked(): if deepseek_info: r = messages_req(token_key, deepseek_info[1]) assert is_blocked(r), f"DeepSeek渠道 + Claude端点 应被拦截" run_test("Claude渠道 + /v1/chat/completions 拦截", test_claude_openai_blocked) run_test("OpenAI渠道 + /v1/messages 拦截", test_openai_claude_blocked) run_test("Gemini渠道 + /v1/chat/completions 拦截", test_gemini_openai_blocked) run_test("OpenAI渠道 + /v1beta 拦截", test_openai_gemini_blocked) run_test("DeepSeek渠道 + /v1/messages 拦截", test_deepseek_claude_blocked) def phase4_error_messages(token_key, openai_model, claude_model, gemini_info): """阶段 4: 错误消息验证""" print("\n--- 阶段 4: 错误消息验证 ---") def test_claude_error_msg(): r = chat_req(token_key, claude_model) assert is_blocked(r) body = r.json() err = body["error"] assert err["code"] == "endpoint_format_restricted" assert err["type"] == "new_api_error" assert "Anthropic" in err["message"] assert "/v1/messages" in err["message"] def test_openai_error_msg(): r = messages_req(token_key, openai_model) assert is_blocked(r) body = r.json() err = body["error"] assert err["code"] == "endpoint_format_restricted" assert "OpenAI" in err["message"] assert "/v1/chat/completions" in err["message"] def test_gemini_error_msg(): if gemini_info: r = chat_req(token_key, gemini_info[1]) assert is_blocked(r) body = r.json() err = body["error"] assert err["code"] == "endpoint_format_restricted" assert "Gemini" in err["message"] assert "/v1beta/models/{model}:generateContent" in err["message"] run_test("Claude渠道错误消息 (Anthropic + /v1/messages)", test_claude_error_msg) run_test("OpenAI渠道错误消息 (OpenAI + /v1/chat/completions)", test_openai_error_msg) run_test("Gemini渠道错误消息 (Gemini + /v1beta)", test_gemini_error_msg) def main(): global BASE_URL, passed, failed parser = argparse.ArgumentParser(description="E2E test: Endpoint Format Restriction") parser.add_argument("--base-url", default=DEFAULT_BASE_URL) parser.add_argument("--db-host", default=DEFAULT_DB_HOST) parser.add_argument("--db-port", type=int, default=DEFAULT_DB_PORT) parser.add_argument("--db-user", default=DEFAULT_DB_USER) parser.add_argument("--db-pass", default=DEFAULT_DB_PASS) parser.add_argument("--db-name", default=DEFAULT_DB_NAME) parser.add_argument("--token-id", type=int, default=29) args = parser.parse_args() BASE_URL = args.base_url print("=" * 60) print("E2E Test: Endpoint Format Restriction") print("=" * 60) # ---- 阶段 0: 环境发现 ---- print("\n--- 阶段 0: 环境发现 ---") # 健康检查 try: r = requests.get(f"{BASE_URL}/api/status", timeout=5) assert r.status_code == 200 print(f" Server: {BASE_URL} OK") except Exception as e: print(f" ERROR: Server not reachable: {e}") sys.exit(1) # 获取 token try: token_key = get_token_key(args) print(f" Token: id={args.token_id} key=sk-{token_key[:8]}...") except ValueError: print(f" ERROR: Token id={args.token_id} not found") sys.exit(1) # 发现渠道 channels = discover_channels(args) openai_ch = find_channel(channels, CH_OPENAI) claude_ch = find_channel(channels, CH_ANTHROPIC) gemini_ch = find_channel(channels, CH_GEMINI) deepseek_ch = find_channel(channels, CH_DEEPSEEK) if not openai_ch: print(" ERROR: No OpenAI channel found") sys.exit(1) if not claude_ch: print(" ERROR: No Anthropic channel found") sys.exit(1) openai_model = openai_ch[1] claude_model = claude_ch[1] gemini_info = gemini_ch # 可能为 None deepseek_info = deepseek_ch # 可能为 None print(f" OpenAI model: {openai_model} (ch={openai_ch[0]})") print(f" Claude model: {claude_model} (ch={claude_ch[0]})") if gemini_info: print(f" Gemini model: {gemini_info[1]} (ch={gemini_info[0]})") else: print(f" Gemini model: (no channel, skip gemini tests)") if deepseek_info: print(f" DeepSeek model: {deepseek_info[1]} (ch={deepseek_info[0]})") # ---- 阶段 1: 功能关闭 ---- set_restriction(args, False) print("\n Set restriction=false, waiting for sync...") wait_for_restriction(args, token_key, claude_model, enabled=False) phase1_off(token_key, openai_model, claude_model, gemini_info) # ---- 阶段 2: 功能开启 - 放行 ---- set_restriction(args, True) print("\n Set restriction=true, waiting for sync...") if not wait_for_restriction(args, token_key, claude_model, enabled=True): print(" ERROR: Restriction did not activate. Aborting.") sys.exit(1) phase2_allowed(token_key, openai_model, claude_model, gemini_info, deepseek_info) # ---- 阶段 3: 功能开启 - 拦截 ---- phase3_blocked(token_key, openai_model, claude_model, gemini_info, deepseek_info) # ---- 阶段 4: 错误消息 ---- phase4_error_messages(token_key, openai_model, claude_model, gemini_info) # ---- 阶段 5: 清理 ---- set_restriction(args, False) print(f"\n--- 阶段 5: 清理 ---") print(f" Restriction reset to: false") # ---- 结果 ---- print("\n" + "=" * 60) print(f"Results: {passed} passed, {failed} failed") if failed > 0: sys.exit(1) else: print("ALL TESTS PASSED!") if __name__ == "__main__": main()