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