Não pode escolher mais do que 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 

592 linhas
22 KiB

  1. package middleware
  2. import (
  3. "encoding/json"
  4. "net/http/httptest"
  5. "sync"
  6. "sync/atomic"
  7. "testing"
  8. "github.com/QuantumNous/new-api/common"
  9. "github.com/QuantumNous/new-api/constant"
  10. "github.com/gin-gonic/gin"
  11. "github.com/stretchr/testify/assert"
  12. "github.com/stretchr/testify/require"
  13. )
  14. // ============================================================================
  15. // Group A: getFamilyFromPath 纯函数测试
  16. // ============================================================================
  17. func TestGetFamilyFromPath_Claude(t *testing.T) {
  18. assert.Equal(t, "claude", getFamilyFromPath("/v1/messages"))
  19. }
  20. func TestGetFamilyFromPath_ClaudeSubpath(t *testing.T) {
  21. assert.Equal(t, "claude", getFamilyFromPath("/v1/messages/"))
  22. }
  23. func TestGetFamilyFromPath_ClaudeBatchPath(t *testing.T) {
  24. assert.Equal(t, "claude", getFamilyFromPath("/v1/messages/batches"))
  25. }
  26. func TestGetFamilyFromPath_Gemini(t *testing.T) {
  27. assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/models/gemini-pro:generateContent"))
  28. assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/"))
  29. }
  30. func TestGetFamilyFromPath_GeminiCompatPath(t *testing.T) {
  31. assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/openai/models"))
  32. }
  33. func TestGetFamilyFromPath_GeminiStreamPath(t *testing.T) {
  34. assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/models/gemini-2.0-flash:streamGenerateContent"))
  35. }
  36. func TestGetFamilyFromPath_V1ModelsPath(t *testing.T) {
  37. // /v1/models/*path 用于 Gemini relay,应归类为 gemini 家族
  38. assert.Equal(t, "gemini", getFamilyFromPath("/v1/models/gemini-pro:generateContent"))
  39. assert.Equal(t, "gemini", getFamilyFromPath("/v1/models/gemini-2.0-flash:streamGenerateContent"))
  40. }
  41. func TestGetFamilyFromPath_V1EnginesPath(t *testing.T) {
  42. // /v1/engines/:model/embeddings 用于 Gemini relay,应归类为 gemini 家族
  43. assert.Equal(t, "gemini", getFamilyFromPath("/v1/engines/text-embedding-004/embeddings"))
  44. }
  45. func TestGetFamilyFromPath_OpenAI(t *testing.T) {
  46. assert.Equal(t, "openai", getFamilyFromPath("/v1/chat/completions"))
  47. assert.Equal(t, "openai", getFamilyFromPath("/v1/completions"))
  48. assert.Equal(t, "openai", getFamilyFromPath("/v1/embeddings"))
  49. assert.Equal(t, "openai", getFamilyFromPath("/v1/audio/speech"))
  50. assert.Equal(t, "openai", getFamilyFromPath("/v1/images/generations"))
  51. assert.Equal(t, "openai", getFamilyFromPath("/v1/rerank"))
  52. assert.Equal(t, "openai", getFamilyFromPath("/v1/moderations"))
  53. }
  54. func TestGetFamilyFromPath_Realtime(t *testing.T) {
  55. assert.Equal(t, "openai", getFamilyFromPath("/v1/realtime"))
  56. }
  57. func TestGetFamilyFromPath_Responses(t *testing.T) {
  58. assert.Equal(t, "openai", getFamilyFromPath("/v1/responses"))
  59. assert.Equal(t, "openai", getFamilyFromPath("/v1/responses/compact"))
  60. }
  61. func TestGetFamilyFromPath_Unknown(t *testing.T) {
  62. assert.Equal(t, "", getFamilyFromPath("/api/status"))
  63. assert.Equal(t, "", getFamilyFromPath("/pg/chat/completions"))
  64. assert.Equal(t, "", getFamilyFromPath("/"))
  65. assert.Equal(t, "", getFamilyFromPath("/v1"))
  66. }
  67. // ============================================================================
  68. // Group B: isChannelAllowed 纯函数测试
  69. // ============================================================================
  70. func TestIsChannelAllowed_OpenAIChannel_OpenAIFormat(t *testing.T) {
  71. assert.True(t, isChannelAllowed(constant.ChannelTypeOpenAI, "openai"))
  72. }
  73. func TestIsChannelAllowed_OpenAIChannel_ClaudeFormat(t *testing.T) {
  74. assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, "claude"))
  75. }
  76. func TestIsChannelAllowed_OpenAIChannel_GeminiFormat(t *testing.T) {
  77. assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, "gemini"))
  78. }
  79. func TestIsChannelAllowed_ClaudeChannel_ClaudeFormat(t *testing.T) {
  80. assert.True(t, isChannelAllowed(constant.ChannelTypeAnthropic, "claude"))
  81. }
  82. func TestIsChannelAllowed_ClaudeChannel_OpenAIFormat(t *testing.T) {
  83. assert.False(t, isChannelAllowed(constant.ChannelTypeAnthropic, "openai"))
  84. }
  85. func TestIsChannelAllowed_GeminiChannel(t *testing.T) {
  86. assert.True(t, isChannelAllowed(constant.ChannelTypeGemini, "gemini"))
  87. assert.False(t, isChannelAllowed(constant.ChannelTypeGemini, "openai"))
  88. assert.False(t, isChannelAllowed(constant.ChannelTypeGemini, "claude"))
  89. }
  90. func TestIsChannelAllowed_AzureChannel(t *testing.T) {
  91. assert.True(t, isChannelAllowed(constant.ChannelTypeAzure, "openai"))
  92. assert.False(t, isChannelAllowed(constant.ChannelTypeAzure, "claude"))
  93. assert.False(t, isChannelAllowed(constant.ChannelTypeAzure, "gemini"))
  94. }
  95. func TestIsChannelAllowed_DeepSeekChannel(t *testing.T) {
  96. assert.True(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "openai"))
  97. assert.False(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "claude"))
  98. assert.False(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "gemini"))
  99. }
  100. func TestIsChannelAllowed_VertexAI_AllFormats(t *testing.T) {
  101. assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "openai"))
  102. assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "claude"))
  103. assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "gemini"))
  104. }
  105. func TestIsChannelAllowed_Ali(t *testing.T) {
  106. assert.True(t, isChannelAllowed(constant.ChannelTypeAli, "openai"))
  107. assert.True(t, isChannelAllowed(constant.ChannelTypeAli, "claude"))
  108. assert.False(t, isChannelAllowed(constant.ChannelTypeAli, "gemini"))
  109. }
  110. func TestIsChannelAllowed_Midjourney_PassThrough(t *testing.T) {
  111. // Midjourney 不在 channelNativeFamilies 中,所有格式都应放行
  112. assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "openai"))
  113. assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "claude"))
  114. assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "gemini"))
  115. }
  116. func TestIsChannelAllowed_Kling_PassThrough(t *testing.T) {
  117. // Kling 不在 channelNativeFamilies 中,所有格式都应放行
  118. assert.True(t, isChannelAllowed(constant.ChannelTypeKling, "openai"))
  119. assert.True(t, isChannelAllowed(constant.ChannelTypeKling, "claude"))
  120. }
  121. func TestIsChannelAllowed_Suno_PassThrough(t *testing.T) {
  122. // Suno 不在 channelNativeFamilies 中
  123. assert.True(t, isChannelAllowed(constant.ChannelTypeSunoAPI, "openai"))
  124. }
  125. func TestIsChannelAllowed_UnknownChannel(t *testing.T) {
  126. assert.True(t, isChannelAllowed(999, "openai"))
  127. assert.True(t, isChannelAllowed(999, "claude"))
  128. assert.True(t, isChannelAllowed(999, "gemini"))
  129. }
  130. func TestIsChannelAllowed_EmptyFamily(t *testing.T) {
  131. // 空 family 表示路径不匹配任何已知格式,guard 会直接放行
  132. // isChannelAllowed 不会被空 family 调用,但验证行为:openai 在 families 中不含 ""
  133. assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, ""))
  134. }
  135. // ============================================================================
  136. // Helper: runGuardTest
  137. // ============================================================================
  138. func runGuardTestWithMethod(method, path string, channelType int) (bool, *httptest.ResponseRecorder) {
  139. called := false
  140. gin.SetMode(gin.TestMode)
  141. r := gin.New()
  142. r.Use(func(c *gin.Context) {
  143. if channelType >= 0 {
  144. c.Set(string(constant.ContextKeyChannelType), channelType)
  145. }
  146. c.Next()
  147. })
  148. r.Use(EndpointFormatGuard())
  149. // 注册所有相关路由
  150. r.POST("/v1/chat/completions", func(ctx *gin.Context) { called = true })
  151. r.POST("/v1/completions", func(ctx *gin.Context) { called = true })
  152. r.POST("/v1/messages", func(ctx *gin.Context) { called = true })
  153. r.POST("/v1/embeddings", func(ctx *gin.Context) { called = true })
  154. r.POST("/v1/rerank", func(ctx *gin.Context) { called = true })
  155. r.POST("/v1/responses", func(ctx *gin.Context) { called = true })
  156. r.POST("/v1/responses/compact", func(ctx *gin.Context) { called = true })
  157. r.POST("/v1/moderations", func(ctx *gin.Context) { called = true })
  158. r.POST("/v1/audio/speech", func(ctx *gin.Context) { called = true })
  159. r.POST("/v1/images/generations", func(ctx *gin.Context) { called = true })
  160. r.POST("/v1/models/*path", func(ctx *gin.Context) { called = true })
  161. r.POST("/v1/engines/:model/embeddings", func(ctx *gin.Context) { called = true })
  162. r.GET("/v1/realtime", func(ctx *gin.Context) { called = true })
  163. r.Any("/v1beta/models/*path", func(ctx *gin.Context) { called = true })
  164. w := httptest.NewRecorder()
  165. req := httptest.NewRequest(method, path, nil)
  166. r.ServeHTTP(w, req)
  167. return called, w
  168. }
  169. func runGuardTest(path string, channelType int) (bool, *httptest.ResponseRecorder) {
  170. return runGuardTestWithMethod("POST", path, channelType)
  171. }
  172. // ============================================================================
  173. // Group C: 中间件集成测试
  174. // ============================================================================
  175. func TestEndpointFormatGuard_Disabled(t *testing.T) {
  176. orig := setTestRestrictionEnabled(false)
  177. defer setTestRestrictionEnabled(orig)
  178. called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
  179. assert.True(t, called, "handler should be called when restriction is disabled")
  180. }
  181. func TestEndpointFormatGuard_Disabled_MultipleChannelTypes(t *testing.T) {
  182. orig := setTestRestrictionEnabled(false)
  183. defer setTestRestrictionEnabled(orig)
  184. for _, ch := range []int{constant.ChannelTypeOpenAI, constant.ChannelTypeAnthropic, constant.ChannelTypeGemini, constant.ChannelTypeMidjourney} {
  185. called, _ := runGuardTest("/v1/chat/completions", ch)
  186. assert.True(t, called, "all channel types should pass when restriction is disabled")
  187. }
  188. }
  189. func TestEndpointFormatGuard_Enabled_Matched(t *testing.T) {
  190. orig := setTestRestrictionEnabled(true)
  191. defer setTestRestrictionEnabled(orig)
  192. called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeOpenAI)
  193. assert.True(t, called, "handler should be called when format matches channel type")
  194. }
  195. func TestEndpointFormatGuard_Enabled_Mismatched(t *testing.T) {
  196. orig := setTestRestrictionEnabled(true)
  197. defer setTestRestrictionEnabled(orig)
  198. called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
  199. assert.False(t, called, "handler should NOT be called when format mismatches")
  200. assert.Equal(t, 400, w.Code)
  201. }
  202. func TestEndpointFormatGuard_Enabled_ClaudeEndpoint_ClaudeChannel(t *testing.T) {
  203. orig := setTestRestrictionEnabled(true)
  204. defer setTestRestrictionEnabled(orig)
  205. called, _ := runGuardTest("/v1/messages", constant.ChannelTypeAnthropic)
  206. assert.True(t, called, "claude endpoint + claude channel should be allowed")
  207. }
  208. func TestEndpointFormatGuard_Enabled_ClaudeEndpoint_OpenAIChannel(t *testing.T) {
  209. orig := setTestRestrictionEnabled(true)
  210. defer setTestRestrictionEnabled(orig)
  211. called, w := runGuardTest("/v1/messages", constant.ChannelTypeOpenAI)
  212. assert.False(t, called, "claude endpoint + openai channel should be blocked")
  213. assert.Equal(t, 400, w.Code)
  214. }
  215. func TestEndpointFormatGuard_Enabled_GeminiEndpoint_GeminiChannel(t *testing.T) {
  216. orig := setTestRestrictionEnabled(true)
  217. defer setTestRestrictionEnabled(orig)
  218. called, _ := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeGemini)
  219. assert.True(t, called, "gemini endpoint + gemini channel should be allowed")
  220. }
  221. func TestEndpointFormatGuard_Enabled_GeminiEndpoint_OpenAIChannel(t *testing.T) {
  222. orig := setTestRestrictionEnabled(true)
  223. defer setTestRestrictionEnabled(orig)
  224. called, w := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeOpenAI)
  225. assert.False(t, called, "gemini endpoint + openai channel should be blocked")
  226. assert.Equal(t, 400, w.Code)
  227. }
  228. func TestEndpointFormatGuard_Enabled_GeminiChannel_OpenAIEndpoint(t *testing.T) {
  229. orig := setTestRestrictionEnabled(true)
  230. defer setTestRestrictionEnabled(orig)
  231. called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeGemini)
  232. assert.False(t, called, "gemini channel should be blocked on openai endpoint")
  233. assert.Equal(t, 400, w.Code)
  234. }
  235. func TestEndpointFormatGuard_Enabled_GeminiChannel_V1ModelsPath(t *testing.T) {
  236. orig := setTestRestrictionEnabled(true)
  237. defer setTestRestrictionEnabled(orig)
  238. called, _ := runGuardTest("/v1/models/gemini-pro:generateContent", constant.ChannelTypeGemini)
  239. assert.True(t, called, "gemini channel should be allowed on /v1/models/* path (gemini family)")
  240. }
  241. func TestEndpointFormatGuard_Enabled_DeepSeekAllowedOnOpenAI(t *testing.T) {
  242. orig := setTestRestrictionEnabled(true)
  243. defer setTestRestrictionEnabled(orig)
  244. called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeDeepSeek)
  245. assert.True(t, called, "deepseek channel should be allowed on openai endpoint")
  246. }
  247. func TestEndpointFormatGuard_Enabled_DeepSeekBlockedOnClaude(t *testing.T) {
  248. orig := setTestRestrictionEnabled(true)
  249. defer setTestRestrictionEnabled(orig)
  250. called, w := runGuardTest("/v1/messages", constant.ChannelTypeDeepSeek)
  251. assert.False(t, called, "deepseek channel should be blocked on claude endpoint")
  252. assert.Equal(t, 400, w.Code)
  253. }
  254. func TestEndpointFormatGuard_Enabled_MidjourneyPassThrough(t *testing.T) {
  255. orig := setTestRestrictionEnabled(true)
  256. defer setTestRestrictionEnabled(orig)
  257. called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeMidjourney)
  258. assert.True(t, called, "midjourney channel should bypass restriction (not LLM)")
  259. }
  260. func TestEndpointFormatGuard_Enabled_KlingPassThrough(t *testing.T) {
  261. orig := setTestRestrictionEnabled(true)
  262. defer setTestRestrictionEnabled(orig)
  263. called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeKling)
  264. assert.True(t, called, "kling channel should bypass restriction (not LLM)")
  265. }
  266. func TestEndpointFormatGuard_Enabled_RealtimeWebSocket(t *testing.T) {
  267. orig := setTestRestrictionEnabled(true)
  268. defer setTestRestrictionEnabled(orig)
  269. called, _ := runGuardTestWithMethod("GET", "/v1/realtime", constant.ChannelTypeOpenAI)
  270. assert.True(t, called, "openai channel should be allowed on realtime websocket")
  271. }
  272. func TestEndpointFormatGuard_Enabled_RealtimeWebSocket_Blocked(t *testing.T) {
  273. orig := setTestRestrictionEnabled(true)
  274. defer setTestRestrictionEnabled(orig)
  275. called, w := runGuardTestWithMethod("GET", "/v1/realtime", constant.ChannelTypeAnthropic)
  276. assert.False(t, called, "anthropic channel should be blocked on realtime websocket")
  277. assert.Equal(t, 400, w.Code)
  278. }
  279. func TestEndpointFormatGuard_Enabled_ResponsesEndpoint(t *testing.T) {
  280. orig := setTestRestrictionEnabled(true)
  281. defer setTestRestrictionEnabled(orig)
  282. called, _ := runGuardTest("/v1/responses", constant.ChannelTypeOpenAI)
  283. assert.True(t, called, "openai channel should be allowed on /v1/responses")
  284. called, w := runGuardTest("/v1/responses", constant.ChannelTypeAnthropic)
  285. assert.False(t, called, "anthropic channel should be blocked on /v1/responses")
  286. assert.Equal(t, 400, w.Code)
  287. }
  288. func TestEndpointFormatGuard_Enabled_AWS_ClaudeAllowed(t *testing.T) {
  289. orig := setTestRestrictionEnabled(true)
  290. defer setTestRestrictionEnabled(orig)
  291. called, _ := runGuardTest("/v1/messages", constant.ChannelTypeAws)
  292. assert.True(t, called, "AWS channel should be allowed for claude endpoint")
  293. }
  294. func TestEndpointFormatGuard_Enabled_AWS_OpenAIBlocked(t *testing.T) {
  295. orig := setTestRestrictionEnabled(true)
  296. defer setTestRestrictionEnabled(orig)
  297. called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAws)
  298. assert.False(t, called, "AWS channel should be blocked for openai endpoint")
  299. assert.Equal(t, 400, w.Code)
  300. }
  301. func TestEndpointFormatGuard_Enabled_VertexAI_AllEndpoints(t *testing.T) {
  302. orig := setTestRestrictionEnabled(true)
  303. defer setTestRestrictionEnabled(orig)
  304. called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeVertexAi)
  305. assert.True(t, called, "vertexai should be allowed on openai endpoint")
  306. called, _ = runGuardTest("/v1/messages", constant.ChannelTypeVertexAi)
  307. assert.True(t, called, "vertexai should be allowed on claude endpoint")
  308. called, _ = runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeVertexAi)
  309. assert.True(t, called, "vertexai should be allowed on gemini endpoint")
  310. }
  311. func TestEndpointFormatGuard_Enabled_Ali_GeminiBlocked(t *testing.T) {
  312. orig := setTestRestrictionEnabled(true)
  313. defer setTestRestrictionEnabled(orig)
  314. called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAli)
  315. assert.True(t, called, "ali should be allowed on openai endpoint")
  316. called, _ = runGuardTest("/v1/messages", constant.ChannelTypeAli)
  317. assert.True(t, called, "ali should be allowed on claude endpoint")
  318. called, w := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeAli)
  319. assert.False(t, called, "ali should be blocked on gemini endpoint")
  320. assert.Equal(t, 400, w.Code)
  321. }
  322. func TestEndpointFormatGuard_Enabled_NoChannelType(t *testing.T) {
  323. orig := setTestRestrictionEnabled(true)
  324. defer setTestRestrictionEnabled(orig)
  325. called, _ := runGuardTest("/v1/chat/completions", -1)
  326. assert.True(t, called, "handler should be called when no channelType in context")
  327. }
  328. func TestEndpointFormatGuard_Enabled_ChannelTypeZero(t *testing.T) {
  329. orig := setTestRestrictionEnabled(true)
  330. defer setTestRestrictionEnabled(orig)
  331. called, _ := runGuardTest("/v1/chat/completions", 0)
  332. assert.True(t, called, "handler should be called when channelType is 0")
  333. }
  334. func TestEndpointFormatGuard_NonV1Path(t *testing.T) {
  335. orig := setTestRestrictionEnabled(true)
  336. defer setTestRestrictionEnabled(orig)
  337. // /api/status 不在 guard 路由组中,无法通过 runGuardTest 测试
  338. // 改用直接调用 getFamilyFromPath 验证返回空
  339. family := getFamilyFromPath("/api/status")
  340. assert.Equal(t, "", family, "non-v1 paths should have empty family")
  341. }
  342. // ============================================================================
  343. // Group D: 错误消息验证
  344. // ============================================================================
  345. func TestEndpointFormatGuard_ErrorMessage_ClaudeChannel(t *testing.T) {
  346. orig := setTestRestrictionEnabled(true)
  347. defer setTestRestrictionEnabled(orig)
  348. _, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
  349. require.Equal(t, 400, w.Code)
  350. var body map[string]interface{}
  351. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  352. errObj := body["error"].(map[string]interface{})
  353. assert.Equal(t, "endpoint_format_restricted", errObj["code"])
  354. assert.Equal(t, "new_api_error", errObj["type"])
  355. msg := errObj["message"].(string)
  356. assert.Contains(t, msg, "Anthropic")
  357. assert.Contains(t, msg, "/v1/messages")
  358. assert.Contains(t, msg, "openai endpoints")
  359. }
  360. func TestEndpointFormatGuard_ErrorMessage_OpenAIChannel(t *testing.T) {
  361. orig := setTestRestrictionEnabled(true)
  362. defer setTestRestrictionEnabled(orig)
  363. _, w := runGuardTest("/v1/messages", constant.ChannelTypeOpenAI)
  364. require.Equal(t, 400, w.Code)
  365. var body map[string]interface{}
  366. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  367. errObj := body["error"].(map[string]interface{})
  368. assert.Equal(t, "endpoint_format_restricted", errObj["code"])
  369. msg := errObj["message"].(string)
  370. assert.Contains(t, msg, "OpenAI")
  371. assert.Contains(t, msg, "/v1/chat/completions")
  372. assert.Contains(t, msg, "claude endpoints")
  373. }
  374. func TestEndpointFormatGuard_ErrorMessage_GeminiChannel(t *testing.T) {
  375. orig := setTestRestrictionEnabled(true)
  376. defer setTestRestrictionEnabled(orig)
  377. _, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeGemini)
  378. require.Equal(t, 400, w.Code)
  379. var body map[string]interface{}
  380. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  381. errObj := body["error"].(map[string]interface{})
  382. assert.Equal(t, "endpoint_format_restricted", errObj["code"])
  383. msg := errObj["message"].(string)
  384. assert.Contains(t, msg, "Gemini")
  385. assert.Contains(t, msg, "/v1beta/models/{model}:generateContent")
  386. assert.Contains(t, msg, "openai endpoints")
  387. }
  388. func TestEndpointFormatGuard_ErrorMessage_DeepSeekChannel(t *testing.T) {
  389. orig := setTestRestrictionEnabled(true)
  390. defer setTestRestrictionEnabled(orig)
  391. _, w := runGuardTest("/v1/messages", constant.ChannelTypeDeepSeek)
  392. require.Equal(t, 400, w.Code)
  393. var body map[string]interface{}
  394. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
  395. errObj := body["error"].(map[string]interface{})
  396. assert.Equal(t, "endpoint_format_restricted", errObj["code"])
  397. msg := errObj["message"].(string)
  398. assert.Contains(t, msg, "DeepSeek")
  399. assert.Contains(t, msg, "/v1/chat/completions")
  400. }
  401. // ============================================================================
  402. // Group E: 并发测试
  403. // ============================================================================
  404. func TestEndpointFormatGuard_ConcurrentRequests(t *testing.T) {
  405. orig := setTestRestrictionEnabled(true)
  406. defer setTestRestrictionEnabled(orig)
  407. var allowedCount, blockedCount int64
  408. var wg sync.WaitGroup
  409. for i := 0; i < 20; i++ {
  410. wg.Add(1)
  411. go func(idx int) {
  412. defer wg.Done()
  413. var ch int
  414. if idx%2 == 0 {
  415. ch = constant.ChannelTypeOpenAI
  416. } else {
  417. ch = constant.ChannelTypeAnthropic
  418. }
  419. called, w := runGuardTest("/v1/chat/completions", ch)
  420. if idx%2 == 0 {
  421. if called && w.Code == 200 {
  422. atomic.AddInt64(&allowedCount, 1)
  423. }
  424. } else {
  425. if !called && w.Code == 400 {
  426. atomic.AddInt64(&blockedCount, 1)
  427. }
  428. }
  429. }(i)
  430. }
  431. wg.Wait()
  432. assert.Equal(t, int64(10), atomic.LoadInt64(&allowedCount), "10 个 OpenAI 请求应放行")
  433. assert.Equal(t, int64(10), atomic.LoadInt64(&blockedCount), "10 个 Anthropic 请求应被拦截")
  434. }
  435. // ============================================================================
  436. // Group F: 功能开关测试
  437. // ============================================================================
  438. func TestEndpointFormatGuard_RuntimeToggle(t *testing.T) {
  439. orig := setTestRestrictionEnabled(false)
  440. defer setTestRestrictionEnabled(orig)
  441. // 阶段 1: 关闭 → 应放行
  442. called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
  443. assert.True(t, called, "should pass when restriction is OFF")
  444. // 阶段 2: 开启 → 应拦截
  445. setTestRestrictionEnabled(true)
  446. called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
  447. assert.False(t, called, "should be blocked when restriction is ON")
  448. assert.Equal(t, 400, w.Code)
  449. // 阶段 3: 再次关闭 → 应放行
  450. setTestRestrictionEnabled(false)
  451. called, _ = runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
  452. assert.True(t, called, "should pass again when restriction is turned OFF")
  453. }
  454. // ============================================================================
  455. // 辅助测试
  456. // ============================================================================
  457. func TestEndpointFormatGuard_ContextKeyInt(t *testing.T) {
  458. gin.SetMode(gin.TestMode)
  459. w := httptest.NewRecorder()
  460. c, _ := gin.CreateTestContext(w)
  461. c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
  462. c.Set(string(constant.ContextKeyChannelType), 14)
  463. val := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
  464. assert.Equal(t, 14, val)
  465. }