You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

421 lines
14 KiB

  1. """E2E test for Endpoint Format Restriction feature.
  2. 完整测试 6 个阶段:
  3. 0. 环境发现(自动查找渠道和模型)
  4. 1. 功能关闭 - 所有跨格式请求不被拦截
  5. 2. 功能开启 - 匹配的渠道+端点放行
  6. 3. 功能开启 - 不匹配的渠道+端点被拦截
  7. 4. 错误消息验证
  8. 5. 清理
  9. """
  10. import argparse
  11. import json
  12. import sys
  13. import time
  14. import mysql.connector
  15. import requests
  16. DEFAULT_BASE_URL = "http://localhost:3000"
  17. DEFAULT_DB_HOST = "localhost"
  18. DEFAULT_DB_PORT = 33306
  19. DEFAULT_DB_USER = "root"
  20. DEFAULT_DB_PASS = "Lanqi123456"
  21. DEFAULT_DB_NAME = "new-api"
  22. OPTION_KEY = "global.endpoint_format_restriction_enabled"
  23. # 渠道类型常量 (与 Go constant 对应)
  24. CH_OPENAI = 1
  25. CH_MIDJOURNEY = 2
  26. CH_AZURE = 3
  27. CH_OLLAMA = 4
  28. CH_ANTHROPIC = 14
  29. CH_GEMINI = 24
  30. CH_DEEPSEEK = 43
  31. CH_VERTEX_AI = 41
  32. CH_ALI = 17
  33. CH_AWS = 33
  34. CH_KLING = 50
  35. CH_SUNO = 36
  36. # 渠道类型 → 默认格式家族
  37. CHANNEL_FAMILIES = {
  38. CH_OPENAI: ["openai"],
  39. CH_AZURE: ["openai"],
  40. CH_DEEPSEEK: ["openai"],
  41. CH_ANTHROPIC: ["claude"],
  42. CH_AWS: ["claude"],
  43. CH_GEMINI: ["gemini"],
  44. CH_VERTEX_AI: ["openai", "claude", "gemini"],
  45. CH_ALI: ["openai", "claude"],
  46. }
  47. # 非LLM渠道(不在 channelNativeFamilies 中)
  48. NON_LLM_CHANNELS = [CH_MIDJOURNEY, CH_KLING, CH_SUNO]
  49. def get_db(args):
  50. return mysql.connector.connect(
  51. host=args.db_host, port=args.db_port,
  52. user=args.db_user, password=args.db_pass, database=args.db_name,
  53. )
  54. def get_token_key(args) -> str:
  55. db = get_db(args)
  56. cur = db.cursor()
  57. cur.execute("SELECT `key` FROM tokens WHERE id = %s", (args.token_id,))
  58. row = cur.fetchone()
  59. cur.close()
  60. db.close()
  61. if not row:
  62. raise ValueError(f"Token id={args.token_id} not found")
  63. return row[0]
  64. def set_restriction(args, enabled: bool):
  65. db = get_db(args)
  66. cur = db.cursor()
  67. val = "true" if enabled else "false"
  68. cur.execute(
  69. "INSERT INTO options (`key`, value) VALUES (%s, %s) ON DUPLICATE KEY UPDATE value = %s",
  70. (OPTION_KEY, val, val),
  71. )
  72. db.commit()
  73. cur.close()
  74. db.close()
  75. def discover_channels(args) -> dict:
  76. """从 DB 查询可用渠道,按 type 分组。返回 {type: [(id, model_name), ...]}"""
  77. db = get_db(args)
  78. cur = db.cursor()
  79. cur.execute("SELECT id, type, models FROM channels WHERE status = 1")
  80. rows = cur.fetchall()
  81. cur.close()
  82. db.close()
  83. result = {}
  84. for ch_id, ch_type, models_str in rows:
  85. if not models_str:
  86. continue
  87. models = [m.strip() for m in models_str.split(",") if m.strip()]
  88. if models:
  89. result.setdefault(ch_type, []).append((ch_id, models[0]))
  90. return result
  91. def find_channel(channels: dict, ch_type: int) -> tuple:
  92. """查找指定类型的渠道,返回 (channel_id, model_name) 或 None"""
  93. entries = channels.get(ch_type)
  94. if entries:
  95. return entries[0]
  96. return None
  97. # ============================================================================
  98. # Request helpers
  99. # ============================================================================
  100. def chat_req(key: str, model: str) -> requests.Response:
  101. return requests.post(
  102. f"{BASE_URL}/v1/chat/completions",
  103. headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"},
  104. json={"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5},
  105. timeout=30,
  106. )
  107. def messages_req(key: str, model: str) -> requests.Response:
  108. return requests.post(
  109. f"{BASE_URL}/v1/messages",
  110. headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"},
  111. json={"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5},
  112. timeout=30,
  113. )
  114. def gemini_req(key: str, model: str) -> requests.Response:
  115. return requests.post(
  116. f"{BASE_URL}/v1beta/models/{model}:generateContent",
  117. headers={"Authorization": f"Bearer sk-{key}", "Content-Type": "application/json"},
  118. json={"contents": [{"parts": [{"text": "hi"}]}]},
  119. timeout=30,
  120. )
  121. def is_blocked(resp: requests.Response) -> bool:
  122. if resp.status_code != 400:
  123. return False
  124. try:
  125. return resp.json().get("error", {}).get("code") == "endpoint_format_restricted"
  126. except Exception:
  127. return False
  128. def wait_for_restriction(args, token_key: str, model: str, enabled: bool) -> bool:
  129. """轮询等待限制设置生效(最长 70s,每 3s 检查一次)"""
  130. for i in range(24):
  131. r = chat_req(token_key, model)
  132. blocked = is_blocked(r)
  133. if enabled and blocked:
  134. print(f" Restriction ACTIVE (after ~{i*3}s)")
  135. return True
  136. if not enabled and not blocked:
  137. print(f" Restriction INACTIVE (after ~{i*3}s)")
  138. return True
  139. time.sleep(3)
  140. print(f" WARNING: Restriction state not confirmed after 70s")
  141. return False
  142. # ============================================================================
  143. # Test runners
  144. # ============================================================================
  145. BASE_URL = DEFAULT_BASE_URL
  146. passed = 0
  147. failed = 0
  148. def run_test(name: str, fn) -> bool:
  149. global passed, failed
  150. try:
  151. fn()
  152. print(f" [PASS] {name}")
  153. passed += 1
  154. return True
  155. except AssertionError as e:
  156. print(f" [FAIL] {name}: {e}")
  157. failed += 1
  158. return False
  159. def phase1_off(token_key, openai_model, claude_model, gemini_info):
  160. """阶段 1: 功能关闭 - 所有跨格式请求不被拦截"""
  161. print("\n--- 阶段 1: 功能关闭 ---")
  162. def test_cross_openai_claude():
  163. r = chat_req(token_key, claude_model)
  164. assert not is_blocked(r), f"Claude模型 + OpenAI端点 不应被拦截, got {r.status_code}: {r.text[:200]}"
  165. def test_cross_claude_openai():
  166. r = messages_req(token_key, openai_model)
  167. assert not is_blocked(r), f"OpenAI模型 + Claude端点 不应被拦截, got {r.status_code}: {r.text[:200]}"
  168. def test_cross_gemini_openai():
  169. if gemini_info:
  170. r = chat_req(token_key, gemini_info[1])
  171. assert not is_blocked(r), f"Gemini模型 + OpenAI端点 不应被拦截"
  172. run_test("Claude模型 + OpenAI端点 不拦截", test_cross_openai_claude)
  173. run_test("OpenAI模型 + Claude端点 不拦截", test_cross_claude_openai)
  174. run_test("Gemini模型 + OpenAI端点 不拦截", test_cross_gemini_openai)
  175. def phase2_allowed(token_key, openai_model, claude_model, gemini_info, deepseek_info):
  176. """阶段 2: 功能开启 - 匹配的渠道+端点放行"""
  177. print("\n--- 阶段 2: 功能开启 - 匹配放行 ---")
  178. def test_openai_openai():
  179. r = chat_req(token_key, openai_model)
  180. assert not is_blocked(r), f"OpenAI渠道 + OpenAI端点 应放行, got {r.status_code}"
  181. def test_claude_claude():
  182. r = messages_req(token_key, claude_model)
  183. assert not is_blocked(r), f"Claude渠道 + Claude端点 应放行, got {r.status_code}"
  184. def test_gemini_gemini():
  185. if gemini_info:
  186. r = gemini_req(token_key, gemini_info[1])
  187. assert not is_blocked(r), f"Gemini渠道 + Gemini端点 应放行"
  188. def test_deepseek_openai():
  189. if deepseek_info:
  190. r = chat_req(token_key, deepseek_info[1])
  191. assert not is_blocked(r), f"DeepSeek渠道 + OpenAI端点 应放行"
  192. run_test("OpenAI渠道 + /v1/chat/completions 放行", test_openai_openai)
  193. run_test("Claude渠道 + /v1/messages 放行", test_claude_claude)
  194. run_test("Gemini渠道 + /v1beta 放行", test_gemini_gemini)
  195. run_test("DeepSeek渠道 + /v1/chat/completions 放行", test_deepseek_openai)
  196. def phase3_blocked(token_key, openai_model, claude_model, gemini_info, deepseek_info):
  197. """阶段 3: 功能开启 - 不匹配的渠道+端点被拦截"""
  198. print("\n--- 阶段 3: 功能开启 - 不匹配拦截 ---")
  199. def test_claude_openai_blocked():
  200. r = chat_req(token_key, claude_model)
  201. assert is_blocked(r), f"Claude渠道 + OpenAI端点 应被拦截, got {r.status_code}: {r.text[:200]}"
  202. def test_openai_claude_blocked():
  203. r = messages_req(token_key, openai_model)
  204. assert is_blocked(r), f"OpenAI渠道 + Claude端点 应被拦截, got {r.status_code}: {r.text[:200]}"
  205. def test_gemini_openai_blocked():
  206. if gemini_info:
  207. r = chat_req(token_key, gemini_info[1])
  208. assert is_blocked(r), f"Gemini渠道 + OpenAI端点 应被拦截"
  209. def test_openai_gemini_blocked():
  210. if gemini_info:
  211. r = gemini_req(token_key, openai_model)
  212. assert is_blocked(r), f"OpenAI渠道 + Gemini端点 应被拦截"
  213. def test_deepseek_claude_blocked():
  214. if deepseek_info:
  215. r = messages_req(token_key, deepseek_info[1])
  216. assert is_blocked(r), f"DeepSeek渠道 + Claude端点 应被拦截"
  217. run_test("Claude渠道 + /v1/chat/completions 拦截", test_claude_openai_blocked)
  218. run_test("OpenAI渠道 + /v1/messages 拦截", test_openai_claude_blocked)
  219. run_test("Gemini渠道 + /v1/chat/completions 拦截", test_gemini_openai_blocked)
  220. run_test("OpenAI渠道 + /v1beta 拦截", test_openai_gemini_blocked)
  221. run_test("DeepSeek渠道 + /v1/messages 拦截", test_deepseek_claude_blocked)
  222. def phase4_error_messages(token_key, openai_model, claude_model, gemini_info):
  223. """阶段 4: 错误消息验证"""
  224. print("\n--- 阶段 4: 错误消息验证 ---")
  225. def test_claude_error_msg():
  226. r = chat_req(token_key, claude_model)
  227. assert is_blocked(r)
  228. body = r.json()
  229. err = body["error"]
  230. assert err["code"] == "endpoint_format_restricted"
  231. assert err["type"] == "new_api_error"
  232. assert "Anthropic" in err["message"]
  233. assert "/v1/messages" in err["message"]
  234. def test_openai_error_msg():
  235. r = messages_req(token_key, openai_model)
  236. assert is_blocked(r)
  237. body = r.json()
  238. err = body["error"]
  239. assert err["code"] == "endpoint_format_restricted"
  240. assert "OpenAI" in err["message"]
  241. assert "/v1/chat/completions" in err["message"]
  242. def test_gemini_error_msg():
  243. if gemini_info:
  244. r = chat_req(token_key, gemini_info[1])
  245. assert is_blocked(r)
  246. body = r.json()
  247. err = body["error"]
  248. assert err["code"] == "endpoint_format_restricted"
  249. assert "Gemini" in err["message"]
  250. assert "/v1beta/models/{model}:generateContent" in err["message"]
  251. run_test("Claude渠道错误消息 (Anthropic + /v1/messages)", test_claude_error_msg)
  252. run_test("OpenAI渠道错误消息 (OpenAI + /v1/chat/completions)", test_openai_error_msg)
  253. run_test("Gemini渠道错误消息 (Gemini + /v1beta)", test_gemini_error_msg)
  254. def main():
  255. global BASE_URL, passed, failed
  256. parser = argparse.ArgumentParser(description="E2E test: Endpoint Format Restriction")
  257. parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
  258. parser.add_argument("--db-host", default=DEFAULT_DB_HOST)
  259. parser.add_argument("--db-port", type=int, default=DEFAULT_DB_PORT)
  260. parser.add_argument("--db-user", default=DEFAULT_DB_USER)
  261. parser.add_argument("--db-pass", default=DEFAULT_DB_PASS)
  262. parser.add_argument("--db-name", default=DEFAULT_DB_NAME)
  263. parser.add_argument("--token-id", type=int, default=29)
  264. args = parser.parse_args()
  265. BASE_URL = args.base_url
  266. print("=" * 60)
  267. print("E2E Test: Endpoint Format Restriction")
  268. print("=" * 60)
  269. # ---- 阶段 0: 环境发现 ----
  270. print("\n--- 阶段 0: 环境发现 ---")
  271. # 健康检查
  272. try:
  273. r = requests.get(f"{BASE_URL}/api/status", timeout=5)
  274. assert r.status_code == 200
  275. print(f" Server: {BASE_URL} OK")
  276. except Exception as e:
  277. print(f" ERROR: Server not reachable: {e}")
  278. sys.exit(1)
  279. # 获取 token
  280. try:
  281. token_key = get_token_key(args)
  282. print(f" Token: id={args.token_id} key=sk-{token_key[:8]}...")
  283. except ValueError:
  284. print(f" ERROR: Token id={args.token_id} not found")
  285. sys.exit(1)
  286. # 发现渠道
  287. channels = discover_channels(args)
  288. openai_ch = find_channel(channels, CH_OPENAI)
  289. claude_ch = find_channel(channels, CH_ANTHROPIC)
  290. gemini_ch = find_channel(channels, CH_GEMINI)
  291. deepseek_ch = find_channel(channels, CH_DEEPSEEK)
  292. if not openai_ch:
  293. print(" ERROR: No OpenAI channel found")
  294. sys.exit(1)
  295. if not claude_ch:
  296. print(" ERROR: No Anthropic channel found")
  297. sys.exit(1)
  298. openai_model = openai_ch[1]
  299. claude_model = claude_ch[1]
  300. gemini_info = gemini_ch # 可能为 None
  301. deepseek_info = deepseek_ch # 可能为 None
  302. print(f" OpenAI model: {openai_model} (ch={openai_ch[0]})")
  303. print(f" Claude model: {claude_model} (ch={claude_ch[0]})")
  304. if gemini_info:
  305. print(f" Gemini model: {gemini_info[1]} (ch={gemini_info[0]})")
  306. else:
  307. print(f" Gemini model: (no channel, skip gemini tests)")
  308. if deepseek_info:
  309. print(f" DeepSeek model: {deepseek_info[1]} (ch={deepseek_info[0]})")
  310. # ---- 阶段 1: 功能关闭 ----
  311. set_restriction(args, False)
  312. print("\n Set restriction=false, waiting for sync...")
  313. wait_for_restriction(args, token_key, claude_model, enabled=False)
  314. phase1_off(token_key, openai_model, claude_model, gemini_info)
  315. # ---- 阶段 2: 功能开启 - 放行 ----
  316. set_restriction(args, True)
  317. print("\n Set restriction=true, waiting for sync...")
  318. if not wait_for_restriction(args, token_key, claude_model, enabled=True):
  319. print(" ERROR: Restriction did not activate. Aborting.")
  320. sys.exit(1)
  321. phase2_allowed(token_key, openai_model, claude_model, gemini_info, deepseek_info)
  322. # ---- 阶段 3: 功能开启 - 拦截 ----
  323. phase3_blocked(token_key, openai_model, claude_model, gemini_info, deepseek_info)
  324. # ---- 阶段 4: 错误消息 ----
  325. phase4_error_messages(token_key, openai_model, claude_model, gemini_info)
  326. # ---- 阶段 5: 清理 ----
  327. set_restriction(args, False)
  328. print(f"\n--- 阶段 5: 清理 ---")
  329. print(f" Restriction reset to: false")
  330. # ---- 结果 ----
  331. print("\n" + "=" * 60)
  332. print(f"Results: {passed} passed, {failed} failed")
  333. if failed > 0:
  334. sys.exit(1)
  335. else:
  336. print("ALL TESTS PASSED!")
  337. if __name__ == "__main__":
  338. main()