|
- """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()
|