package middleware import ( "bytes" "fmt" "net/http" "net/http/httptest" "os" "path/filepath" "sync/atomic" "testing" "time" "github.com/QuantumNous/new-api/common" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func setupCaptureTestRouter() *gin.Engine { gin.SetMode(gin.TestMode) return gin.New() } // captureTestCaptureDir 集成测试共享的捕获文件目录 const captureTestDir = "./data/relay-capture" // --- 非流式请求测试 --- func TestRelayCapture_NonStreaming_EnabledUser(t *testing.T) { clearCaptureCache() SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 42) // must be int, not int64 — matches production TokenAuth behavior c.Set(common.RequestIdKey, "test-non-stream-"+time.Now().Format("150405")) c.Next() }) router.Use(RelayCaptureMiddleware()) router.POST("/v1/chat/completions", func(c *gin.Context) { c.JSON(200, gin.H{ "id": "chatcmpl-test", "object": "chat.completion", "choices": []gin.H{{"message": gin.H{"role": "assistant", "content": "Hello!"}}}, }) }) body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}` req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body))) req.Header.Set("Authorization", "Bearer sk-test") req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() router.ServeHTTP(rec, req) assert.Equal(t, 200, rec.Code) today := time.Now().Format("2006-01-02") dateDir := filepath.Join(captureTestDir, today) files, err := os.ReadDir(dateDir) require.NoError(t, err) require.True(t, len(files) > 0, "should have capture log files") logPath := filepath.Join(dateDir, files[len(files)-1].Name()) content, err := os.ReadFile(logPath) require.NoError(t, err) contentStr := string(content) t.Logf("Capture file content:\n%s", contentStr) assert.Contains(t, contentStr, "=== REQUEST") assert.Contains(t, contentStr, "POST /v1/chat/completions") assert.Contains(t, contentStr, "user_id: 42") assert.Contains(t, contentStr, "Authorization: Bearer sk-test") assert.Contains(t, contentStr, "=== RESPONSE") assert.Contains(t, contentStr, "chatcmpl-test") assert.Contains(t, contentStr, "=== END") assert.Contains(t, contentStr, "duration_ms=") } // --- 流式 SSE 请求测试 --- func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { clearCaptureCache() t.Cleanup(func() { os.RemoveAll(captureTestDir) }) SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 42) c.Set(common.RequestIdKey, "test-stream-"+time.Now().Format("150405")) c.Next() }) router.Use(RelayCaptureMiddleware()) router.POST("/v1/chat/completions", func(c *gin.Context) { c.Header("Content-Type", "text/event-stream") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") flusher := c.Writer.(http.Flusher) chunks := []string{ "data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n", "data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"world!\"}}]}\n\n", "data: [DONE]\n\n", } for _, chunk := range chunks { c.Writer.Write([]byte(chunk)) flusher.Flush() time.Sleep(5 * time.Millisecond) } }) body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}],"stream":true}` req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body))) req.Header.Set("Authorization", "Bearer sk-test") req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() router.ServeHTTP(rec, req) assert.Equal(t, 200, rec.Code) today := time.Now().Format("2006-01-02") dateDir := filepath.Join(captureTestDir, today) files, err := os.ReadDir(dateDir) require.NoError(t, err) require.True(t, len(files) > 0, "should have capture log files") logPath := filepath.Join(dateDir, files[len(files)-1].Name()) content, err := os.ReadFile(logPath) require.NoError(t, err) contentStr := string(content) t.Logf("Stream capture file content:\n%s", contentStr) assert.Contains(t, contentStr, "=== REQUEST") assert.Contains(t, contentStr, "POST /v1/chat/completions") assert.Contains(t, contentStr, "user_id: 42") assert.Contains(t, contentStr, "=== RESPONSE") assert.Contains(t, contentStr, "Hello ") assert.Contains(t, contentStr, "world!") assert.Contains(t, contentStr, "[DONE]") assert.Contains(t, contentStr, "=== END") assert.Contains(t, contentStr, "duration_ms=") } // --- 未开启抓包用户测试 --- func TestRelayCapture_DisabledUser(t *testing.T) { clearCaptureCache() router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 99) c.Set(common.RequestIdKey, "test-disabled-"+time.Now().Format("150405")) c.Next() }) router.Use(RelayCaptureMiddleware()) router.POST("/v1/chat/completions", func(c *gin.Context) { c.JSON(200, gin.H{"result": "ok"}) }) req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{}`))) req.Header.Set("Authorization", "Bearer sk-test") rec := httptest.NewRecorder() router.ServeHTTP(rec, req) assert.Equal(t, 200, rec.Code) today := time.Now().Format("2006-01-02") dateDir := filepath.Join(captureTestDir, today) if _, err := os.Stat(dateDir); err == nil { files, _ := os.ReadDir(dateDir) for _, f := range files { assert.NotContains(t, f.Name(), "test-disabled", "disabled user should not generate capture files") } } } // --- 大量数据流测试 --- func TestRelayCapture_LargeStream(t *testing.T) { clearCaptureCache() t.Cleanup(func() { os.RemoveAll(captureTestDir) }) SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 42) c.Set(common.RequestIdKey, "test-large-"+time.Now().Format("150405")) c.Next() }) router.Use(RelayCaptureMiddleware()) router.POST("/v1/chat/completions", func(c *gin.Context) { c.Header("Content-Type", "text/event-stream") flusher := c.Writer.(http.Flusher) for i := 0; i < 50; i++ { chunk := fmt.Sprintf("data: {\"id\":\"%d\",\"choices\":[{\"delta\":{\"content\":\"chunk%d \"}}]}\n\n", i, i) c.Writer.Write([]byte(chunk)) flusher.Flush() } c.Writer.Write([]byte("data: [DONE]\n\n")) flusher.Flush() }) req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4","stream":true}`))) req.Header.Set("Authorization", "Bearer sk-test") rec := httptest.NewRecorder() router.ServeHTTP(rec, req) assert.Equal(t, 200, rec.Code) today := time.Now().Format("2006-01-02") dateDir := filepath.Join(captureTestDir, today) files, err := os.ReadDir(dateDir) require.NoError(t, err) require.True(t, len(files) > 0) logPath := filepath.Join(dateDir, files[len(files)-1].Name()) content, err := os.ReadFile(logPath) require.NoError(t, err) contentStr := string(content) assert.Contains(t, contentStr, "chunk0") assert.Contains(t, contentStr, "chunk49") assert.Contains(t, contentStr, "[DONE]") } // --- GET 请求测试 --- func TestRelayCapture_GetRequest(t *testing.T) { clearCaptureCache() t.Cleanup(func() { os.RemoveAll(captureTestDir) }) SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 42) c.Set(common.RequestIdKey, "test-get-"+time.Now().Format("150405")) c.Next() }) router.Use(RelayCaptureMiddleware()) router.GET("/v1/models", func(c *gin.Context) { c.JSON(200, gin.H{ "data": []gin.H{{"id": "gpt-4", "object": "model"}}, }) }) req, _ := http.NewRequest("GET", "/v1/models", nil) req.Header.Set("Authorization", "Bearer sk-test") rec := httptest.NewRecorder() router.ServeHTTP(rec, req) assert.Equal(t, 200, rec.Code) today := time.Now().Format("2006-01-02") dateDir := filepath.Join(captureTestDir, today) files, err := os.ReadDir(dateDir) require.NoError(t, err) require.True(t, len(files) > 0) logPath := filepath.Join(dateDir, files[len(files)-1].Name()) content, err := os.ReadFile(logPath) require.NoError(t, err) contentStr := string(content) assert.Contains(t, contentStr, "=== REQUEST") assert.Contains(t, contentStr, "GET /v1/models") assert.Contains(t, contentStr, "=== RESPONSE") assert.Contains(t, contentStr, "gpt-4") assert.Contains(t, contentStr, "=== END") } // --- 零 ID 用户测试 --- func TestRelayCapture_ZeroUserID(t *testing.T) { clearCaptureCache() router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set(common.RequestIdKey, "test-zero-"+time.Now().Format("150405")) c.Next() }) router.Use(RelayCaptureMiddleware()) router.GET("/v1/models", func(c *gin.Context) { c.JSON(200, gin.H{"data": []gin.H{}}) }) req, _ := http.NewRequest("GET", "/v1/models", nil) rec := httptest.NewRecorder() router.ServeHTTP(rec, req) assert.Equal(t, 200, rec.Code) } // --- 验证响应不被篡改 --- func TestRelayCapture_ResponseUnchanged(t *testing.T) { clearCaptureCache() t.Cleanup(func() { os.RemoveAll(captureTestDir) }) SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 42) c.Set(common.RequestIdKey, "test-unchanged-"+time.Now().Format("150405")) c.Next() }) router.Use(RelayCaptureMiddleware()) router.POST("/v1/chat/completions", func(c *gin.Context) { c.JSON(200, gin.H{ "id": "chatcmpl-123", "choices": []gin.H{{"message": gin.H{"content": "exact response"}}}, }) }) req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4"}`))) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() router.ServeHTTP(rec, req) assert.Equal(t, 200, rec.Code) assert.Contains(t, rec.Body.String(), "exact response") assert.Contains(t, rec.Body.String(), "chatcmpl-123") } // --- 并发安全测试 --- func TestRelayCapture_ConcurrentRequests(t *testing.T) { clearCaptureCache() t.Cleanup(func() { os.RemoveAll(captureTestDir) }) SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) router := setupCaptureTestRouter() var reqCounter int64 = 0 router.Use(func(c *gin.Context) { c.Set("id", 42) n := atomic.AddInt64(&reqCounter, 1) c.Set(common.RequestIdKey, fmt.Sprintf("concurrent-%d-%d", n, time.Now().UnixNano())) c.Next() }) router.Use(RelayCaptureMiddleware()) router.POST("/v1/chat/completions", func(c *gin.Context) { c.JSON(200, gin.H{"status": "ok"}) }) done := make(chan bool, 10) for i := 0; i < 10; i++ { go func(idx int) { req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(fmt.Sprintf(`{"req":%d}`, idx)))) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() router.ServeHTTP(rec, req) assert.Equal(t, 200, rec.Code) done <- true }(i) } for i := 0; i < 10; i++ { select { case <-done: case <-time.After(10 * time.Second): t.Fatal("timeout waiting for concurrent request") } } time.Sleep(100 * time.Millisecond) today := time.Now().Format("2006-01-02") dateDir := filepath.Join(captureTestDir, today) files, err := os.ReadDir(dateDir) require.NoError(t, err) assert.Equal(t, 10, len(files), "should have 10 capture log files") for _, f := range files { content, err := os.ReadFile(filepath.Join(dateDir, f.Name())) require.NoError(t, err) contentStr := string(content) assert.Contains(t, contentStr, "=== REQUEST") assert.Contains(t, contentStr, "=== RESPONSE") assert.Contains(t, contentStr, "=== END") } }