|
- package middleware
-
- import (
- "encoding/json"
- "net/http/httptest"
- "sync"
- "sync/atomic"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/constant"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- )
-
- // ============================================================================
- // Group A: getFamilyFromPath 纯函数测试
- // ============================================================================
-
- func TestGetFamilyFromPath_Claude(t *testing.T) {
- assert.Equal(t, "claude", getFamilyFromPath("/v1/messages"))
- }
-
- func TestGetFamilyFromPath_ClaudeSubpath(t *testing.T) {
- assert.Equal(t, "claude", getFamilyFromPath("/v1/messages/"))
- }
-
- func TestGetFamilyFromPath_ClaudeBatchPath(t *testing.T) {
- assert.Equal(t, "claude", getFamilyFromPath("/v1/messages/batches"))
- }
-
- func TestGetFamilyFromPath_Gemini(t *testing.T) {
- assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/models/gemini-pro:generateContent"))
- assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/"))
- }
-
- func TestGetFamilyFromPath_GeminiCompatPath(t *testing.T) {
- assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/openai/models"))
- }
-
- func TestGetFamilyFromPath_GeminiStreamPath(t *testing.T) {
- assert.Equal(t, "gemini", getFamilyFromPath("/v1beta/models/gemini-2.0-flash:streamGenerateContent"))
- }
-
- func TestGetFamilyFromPath_V1ModelsPath(t *testing.T) {
- // /v1/models/*path 用于 Gemini relay,应归类为 gemini 家族
- assert.Equal(t, "gemini", getFamilyFromPath("/v1/models/gemini-pro:generateContent"))
- assert.Equal(t, "gemini", getFamilyFromPath("/v1/models/gemini-2.0-flash:streamGenerateContent"))
- }
-
- func TestGetFamilyFromPath_V1EnginesPath(t *testing.T) {
- // /v1/engines/:model/embeddings 用于 Gemini relay,应归类为 gemini 家族
- assert.Equal(t, "gemini", getFamilyFromPath("/v1/engines/text-embedding-004/embeddings"))
- }
-
- func TestGetFamilyFromPath_OpenAI(t *testing.T) {
- assert.Equal(t, "openai", getFamilyFromPath("/v1/chat/completions"))
- assert.Equal(t, "openai", getFamilyFromPath("/v1/completions"))
- assert.Equal(t, "openai", getFamilyFromPath("/v1/embeddings"))
- assert.Equal(t, "openai", getFamilyFromPath("/v1/audio/speech"))
- assert.Equal(t, "openai", getFamilyFromPath("/v1/images/generations"))
- assert.Equal(t, "openai", getFamilyFromPath("/v1/rerank"))
- assert.Equal(t, "openai", getFamilyFromPath("/v1/moderations"))
- }
-
- func TestGetFamilyFromPath_Realtime(t *testing.T) {
- assert.Equal(t, "openai", getFamilyFromPath("/v1/realtime"))
- }
-
- func TestGetFamilyFromPath_Responses(t *testing.T) {
- assert.Equal(t, "openai", getFamilyFromPath("/v1/responses"))
- assert.Equal(t, "openai", getFamilyFromPath("/v1/responses/compact"))
- }
-
- func TestGetFamilyFromPath_Unknown(t *testing.T) {
- assert.Equal(t, "", getFamilyFromPath("/api/status"))
- assert.Equal(t, "", getFamilyFromPath("/pg/chat/completions"))
- assert.Equal(t, "", getFamilyFromPath("/"))
- assert.Equal(t, "", getFamilyFromPath("/v1"))
- }
-
- // ============================================================================
- // Group B: isChannelAllowed 纯函数测试
- // ============================================================================
-
- func TestIsChannelAllowed_OpenAIChannel_OpenAIFormat(t *testing.T) {
- assert.True(t, isChannelAllowed(constant.ChannelTypeOpenAI, "openai"))
- }
-
- func TestIsChannelAllowed_OpenAIChannel_ClaudeFormat(t *testing.T) {
- assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, "claude"))
- }
-
- func TestIsChannelAllowed_OpenAIChannel_GeminiFormat(t *testing.T) {
- assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, "gemini"))
- }
-
- func TestIsChannelAllowed_ClaudeChannel_ClaudeFormat(t *testing.T) {
- assert.True(t, isChannelAllowed(constant.ChannelTypeAnthropic, "claude"))
- }
-
- func TestIsChannelAllowed_ClaudeChannel_OpenAIFormat(t *testing.T) {
- assert.False(t, isChannelAllowed(constant.ChannelTypeAnthropic, "openai"))
- }
-
- func TestIsChannelAllowed_GeminiChannel(t *testing.T) {
- assert.True(t, isChannelAllowed(constant.ChannelTypeGemini, "gemini"))
- assert.False(t, isChannelAllowed(constant.ChannelTypeGemini, "openai"))
- assert.False(t, isChannelAllowed(constant.ChannelTypeGemini, "claude"))
- }
-
- func TestIsChannelAllowed_AzureChannel(t *testing.T) {
- assert.True(t, isChannelAllowed(constant.ChannelTypeAzure, "openai"))
- assert.False(t, isChannelAllowed(constant.ChannelTypeAzure, "claude"))
- assert.False(t, isChannelAllowed(constant.ChannelTypeAzure, "gemini"))
- }
-
- func TestIsChannelAllowed_DeepSeekChannel(t *testing.T) {
- assert.True(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "openai"))
- assert.False(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "claude"))
- assert.False(t, isChannelAllowed(constant.ChannelTypeDeepSeek, "gemini"))
- }
-
- func TestIsChannelAllowed_VertexAI_AllFormats(t *testing.T) {
- assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "openai"))
- assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "claude"))
- assert.True(t, isChannelAllowed(constant.ChannelTypeVertexAi, "gemini"))
- }
-
- func TestIsChannelAllowed_Ali(t *testing.T) {
- assert.True(t, isChannelAllowed(constant.ChannelTypeAli, "openai"))
- assert.True(t, isChannelAllowed(constant.ChannelTypeAli, "claude"))
- assert.False(t, isChannelAllowed(constant.ChannelTypeAli, "gemini"))
- }
-
- func TestIsChannelAllowed_Midjourney_PassThrough(t *testing.T) {
- // Midjourney 不在 channelNativeFamilies 中,所有格式都应放行
- assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "openai"))
- assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "claude"))
- assert.True(t, isChannelAllowed(constant.ChannelTypeMidjourney, "gemini"))
- }
-
- func TestIsChannelAllowed_Kling_PassThrough(t *testing.T) {
- // Kling 不在 channelNativeFamilies 中,所有格式都应放行
- assert.True(t, isChannelAllowed(constant.ChannelTypeKling, "openai"))
- assert.True(t, isChannelAllowed(constant.ChannelTypeKling, "claude"))
- }
-
- func TestIsChannelAllowed_Suno_PassThrough(t *testing.T) {
- // Suno 不在 channelNativeFamilies 中
- assert.True(t, isChannelAllowed(constant.ChannelTypeSunoAPI, "openai"))
- }
-
- func TestIsChannelAllowed_UnknownChannel(t *testing.T) {
- assert.True(t, isChannelAllowed(999, "openai"))
- assert.True(t, isChannelAllowed(999, "claude"))
- assert.True(t, isChannelAllowed(999, "gemini"))
- }
-
- func TestIsChannelAllowed_EmptyFamily(t *testing.T) {
- // 空 family 表示路径不匹配任何已知格式,guard 会直接放行
- // isChannelAllowed 不会被空 family 调用,但验证行为:openai 在 families 中不含 ""
- assert.False(t, isChannelAllowed(constant.ChannelTypeOpenAI, ""))
- }
-
- // ============================================================================
- // Helper: runGuardTest
- // ============================================================================
-
- func runGuardTestWithMethod(method, path string, channelType int) (bool, *httptest.ResponseRecorder) {
- called := false
- gin.SetMode(gin.TestMode)
-
- r := gin.New()
- r.Use(func(c *gin.Context) {
- if channelType >= 0 {
- c.Set(string(constant.ContextKeyChannelType), channelType)
- }
- c.Next()
- })
- r.Use(EndpointFormatGuard())
-
- // 注册所有相关路由
- r.POST("/v1/chat/completions", func(ctx *gin.Context) { called = true })
- r.POST("/v1/completions", func(ctx *gin.Context) { called = true })
- r.POST("/v1/messages", func(ctx *gin.Context) { called = true })
- r.POST("/v1/embeddings", func(ctx *gin.Context) { called = true })
- r.POST("/v1/rerank", func(ctx *gin.Context) { called = true })
- r.POST("/v1/responses", func(ctx *gin.Context) { called = true })
- r.POST("/v1/responses/compact", func(ctx *gin.Context) { called = true })
- r.POST("/v1/moderations", func(ctx *gin.Context) { called = true })
- r.POST("/v1/audio/speech", func(ctx *gin.Context) { called = true })
- r.POST("/v1/images/generations", func(ctx *gin.Context) { called = true })
- r.POST("/v1/models/*path", func(ctx *gin.Context) { called = true })
- r.POST("/v1/engines/:model/embeddings", func(ctx *gin.Context) { called = true })
- r.GET("/v1/realtime", func(ctx *gin.Context) { called = true })
- r.Any("/v1beta/models/*path", func(ctx *gin.Context) { called = true })
-
- w := httptest.NewRecorder()
- req := httptest.NewRequest(method, path, nil)
- r.ServeHTTP(w, req)
-
- return called, w
- }
-
- func runGuardTest(path string, channelType int) (bool, *httptest.ResponseRecorder) {
- return runGuardTestWithMethod("POST", path, channelType)
- }
-
- // ============================================================================
- // Group C: 中间件集成测试
- // ============================================================================
-
- func TestEndpointFormatGuard_Disabled(t *testing.T) {
- orig := setTestRestrictionEnabled(false)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
- assert.True(t, called, "handler should be called when restriction is disabled")
- }
-
- func TestEndpointFormatGuard_Disabled_MultipleChannelTypes(t *testing.T) {
- orig := setTestRestrictionEnabled(false)
- defer setTestRestrictionEnabled(orig)
-
- for _, ch := range []int{constant.ChannelTypeOpenAI, constant.ChannelTypeAnthropic, constant.ChannelTypeGemini, constant.ChannelTypeMidjourney} {
- called, _ := runGuardTest("/v1/chat/completions", ch)
- assert.True(t, called, "all channel types should pass when restriction is disabled")
- }
- }
-
- func TestEndpointFormatGuard_Enabled_Matched(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeOpenAI)
- assert.True(t, called, "handler should be called when format matches channel type")
- }
-
- func TestEndpointFormatGuard_Enabled_Mismatched(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
- assert.False(t, called, "handler should NOT be called when format mismatches")
- assert.Equal(t, 400, w.Code)
- }
-
- func TestEndpointFormatGuard_Enabled_ClaudeEndpoint_ClaudeChannel(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/messages", constant.ChannelTypeAnthropic)
- assert.True(t, called, "claude endpoint + claude channel should be allowed")
- }
-
- func TestEndpointFormatGuard_Enabled_ClaudeEndpoint_OpenAIChannel(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, w := runGuardTest("/v1/messages", constant.ChannelTypeOpenAI)
- assert.False(t, called, "claude endpoint + openai channel should be blocked")
- assert.Equal(t, 400, w.Code)
- }
-
- func TestEndpointFormatGuard_Enabled_GeminiEndpoint_GeminiChannel(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeGemini)
- assert.True(t, called, "gemini endpoint + gemini channel should be allowed")
- }
-
- func TestEndpointFormatGuard_Enabled_GeminiEndpoint_OpenAIChannel(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, w := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeOpenAI)
- assert.False(t, called, "gemini endpoint + openai channel should be blocked")
- assert.Equal(t, 400, w.Code)
- }
-
- func TestEndpointFormatGuard_Enabled_GeminiChannel_OpenAIEndpoint(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeGemini)
- assert.False(t, called, "gemini channel should be blocked on openai endpoint")
- assert.Equal(t, 400, w.Code)
- }
-
- func TestEndpointFormatGuard_Enabled_GeminiChannel_V1ModelsPath(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/models/gemini-pro:generateContent", constant.ChannelTypeGemini)
- assert.True(t, called, "gemini channel should be allowed on /v1/models/* path (gemini family)")
- }
-
- func TestEndpointFormatGuard_Enabled_DeepSeekAllowedOnOpenAI(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeDeepSeek)
- assert.True(t, called, "deepseek channel should be allowed on openai endpoint")
- }
-
- func TestEndpointFormatGuard_Enabled_DeepSeekBlockedOnClaude(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, w := runGuardTest("/v1/messages", constant.ChannelTypeDeepSeek)
- assert.False(t, called, "deepseek channel should be blocked on claude endpoint")
- assert.Equal(t, 400, w.Code)
- }
-
- func TestEndpointFormatGuard_Enabled_MidjourneyPassThrough(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeMidjourney)
- assert.True(t, called, "midjourney channel should bypass restriction (not LLM)")
- }
-
- func TestEndpointFormatGuard_Enabled_KlingPassThrough(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeKling)
- assert.True(t, called, "kling channel should bypass restriction (not LLM)")
- }
-
- func TestEndpointFormatGuard_Enabled_RealtimeWebSocket(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTestWithMethod("GET", "/v1/realtime", constant.ChannelTypeOpenAI)
- assert.True(t, called, "openai channel should be allowed on realtime websocket")
- }
-
- func TestEndpointFormatGuard_Enabled_RealtimeWebSocket_Blocked(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, w := runGuardTestWithMethod("GET", "/v1/realtime", constant.ChannelTypeAnthropic)
- assert.False(t, called, "anthropic channel should be blocked on realtime websocket")
- assert.Equal(t, 400, w.Code)
- }
-
- func TestEndpointFormatGuard_Enabled_ResponsesEndpoint(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/responses", constant.ChannelTypeOpenAI)
- assert.True(t, called, "openai channel should be allowed on /v1/responses")
-
- called, w := runGuardTest("/v1/responses", constant.ChannelTypeAnthropic)
- assert.False(t, called, "anthropic channel should be blocked on /v1/responses")
- assert.Equal(t, 400, w.Code)
- }
-
- func TestEndpointFormatGuard_Enabled_AWS_ClaudeAllowed(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/messages", constant.ChannelTypeAws)
- assert.True(t, called, "AWS channel should be allowed for claude endpoint")
- }
-
- func TestEndpointFormatGuard_Enabled_AWS_OpenAIBlocked(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAws)
- assert.False(t, called, "AWS channel should be blocked for openai endpoint")
- assert.Equal(t, 400, w.Code)
- }
-
- func TestEndpointFormatGuard_Enabled_VertexAI_AllEndpoints(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeVertexAi)
- assert.True(t, called, "vertexai should be allowed on openai endpoint")
-
- called, _ = runGuardTest("/v1/messages", constant.ChannelTypeVertexAi)
- assert.True(t, called, "vertexai should be allowed on claude endpoint")
-
- called, _ = runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeVertexAi)
- assert.True(t, called, "vertexai should be allowed on gemini endpoint")
- }
-
- func TestEndpointFormatGuard_Enabled_Ali_GeminiBlocked(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAli)
- assert.True(t, called, "ali should be allowed on openai endpoint")
-
- called, _ = runGuardTest("/v1/messages", constant.ChannelTypeAli)
- assert.True(t, called, "ali should be allowed on claude endpoint")
-
- called, w := runGuardTest("/v1beta/models/gemini-pro:generateContent", constant.ChannelTypeAli)
- assert.False(t, called, "ali should be blocked on gemini endpoint")
- assert.Equal(t, 400, w.Code)
- }
-
- func TestEndpointFormatGuard_Enabled_NoChannelType(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/chat/completions", -1)
- assert.True(t, called, "handler should be called when no channelType in context")
- }
-
- func TestEndpointFormatGuard_Enabled_ChannelTypeZero(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- called, _ := runGuardTest("/v1/chat/completions", 0)
- assert.True(t, called, "handler should be called when channelType is 0")
- }
-
- func TestEndpointFormatGuard_NonV1Path(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- // /api/status 不在 guard 路由组中,无法通过 runGuardTest 测试
- // 改用直接调用 getFamilyFromPath 验证返回空
- family := getFamilyFromPath("/api/status")
- assert.Equal(t, "", family, "non-v1 paths should have empty family")
- }
-
- // ============================================================================
- // Group D: 错误消息验证
- // ============================================================================
-
- func TestEndpointFormatGuard_ErrorMessage_ClaudeChannel(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- _, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
- require.Equal(t, 400, w.Code)
-
- var body map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
-
- errObj := body["error"].(map[string]interface{})
- assert.Equal(t, "endpoint_format_restricted", errObj["code"])
- assert.Equal(t, "new_api_error", errObj["type"])
-
- msg := errObj["message"].(string)
- assert.Contains(t, msg, "Anthropic")
- assert.Contains(t, msg, "/v1/messages")
- assert.Contains(t, msg, "openai endpoints")
- }
-
- func TestEndpointFormatGuard_ErrorMessage_OpenAIChannel(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- _, w := runGuardTest("/v1/messages", constant.ChannelTypeOpenAI)
- require.Equal(t, 400, w.Code)
-
- var body map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
-
- errObj := body["error"].(map[string]interface{})
- assert.Equal(t, "endpoint_format_restricted", errObj["code"])
-
- msg := errObj["message"].(string)
- assert.Contains(t, msg, "OpenAI")
- assert.Contains(t, msg, "/v1/chat/completions")
- assert.Contains(t, msg, "claude endpoints")
- }
-
- func TestEndpointFormatGuard_ErrorMessage_GeminiChannel(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- _, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeGemini)
- require.Equal(t, 400, w.Code)
-
- var body map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
-
- errObj := body["error"].(map[string]interface{})
- assert.Equal(t, "endpoint_format_restricted", errObj["code"])
-
- msg := errObj["message"].(string)
- assert.Contains(t, msg, "Gemini")
- assert.Contains(t, msg, "/v1beta/models/{model}:generateContent")
- assert.Contains(t, msg, "openai endpoints")
- }
-
- func TestEndpointFormatGuard_ErrorMessage_DeepSeekChannel(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- _, w := runGuardTest("/v1/messages", constant.ChannelTypeDeepSeek)
- require.Equal(t, 400, w.Code)
-
- var body map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
-
- errObj := body["error"].(map[string]interface{})
- assert.Equal(t, "endpoint_format_restricted", errObj["code"])
-
- msg := errObj["message"].(string)
- assert.Contains(t, msg, "DeepSeek")
- assert.Contains(t, msg, "/v1/chat/completions")
- }
-
- // ============================================================================
- // Group E: 并发测试
- // ============================================================================
-
- func TestEndpointFormatGuard_ConcurrentRequests(t *testing.T) {
- orig := setTestRestrictionEnabled(true)
- defer setTestRestrictionEnabled(orig)
-
- var allowedCount, blockedCount int64
- var wg sync.WaitGroup
-
- for i := 0; i < 20; i++ {
- wg.Add(1)
- go func(idx int) {
- defer wg.Done()
- var ch int
- if idx%2 == 0 {
- ch = constant.ChannelTypeOpenAI
- } else {
- ch = constant.ChannelTypeAnthropic
- }
- called, w := runGuardTest("/v1/chat/completions", ch)
- if idx%2 == 0 {
- if called && w.Code == 200 {
- atomic.AddInt64(&allowedCount, 1)
- }
- } else {
- if !called && w.Code == 400 {
- atomic.AddInt64(&blockedCount, 1)
- }
- }
- }(i)
- }
- wg.Wait()
-
- assert.Equal(t, int64(10), atomic.LoadInt64(&allowedCount), "10 个 OpenAI 请求应放行")
- assert.Equal(t, int64(10), atomic.LoadInt64(&blockedCount), "10 个 Anthropic 请求应被拦截")
- }
-
- // ============================================================================
- // Group F: 功能开关测试
- // ============================================================================
-
- func TestEndpointFormatGuard_RuntimeToggle(t *testing.T) {
- orig := setTestRestrictionEnabled(false)
- defer setTestRestrictionEnabled(orig)
-
- // 阶段 1: 关闭 → 应放行
- called, _ := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
- assert.True(t, called, "should pass when restriction is OFF")
-
- // 阶段 2: 开启 → 应拦截
- setTestRestrictionEnabled(true)
- called, w := runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
- assert.False(t, called, "should be blocked when restriction is ON")
- assert.Equal(t, 400, w.Code)
-
- // 阶段 3: 再次关闭 → 应放行
- setTestRestrictionEnabled(false)
- called, _ = runGuardTest("/v1/chat/completions", constant.ChannelTypeAnthropic)
- assert.True(t, called, "should pass again when restriction is turned OFF")
- }
-
- // ============================================================================
- // 辅助测试
- // ============================================================================
-
- func TestEndpointFormatGuard_ContextKeyInt(t *testing.T) {
- gin.SetMode(gin.TestMode)
- w := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(w)
- c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
- c.Set(string(constant.ContextKeyChannelType), 14)
-
- val := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
- assert.Equal(t, 14, val)
- }
|