Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

403 строки
12 KiB

  1. package middleware
  2. import (
  3. "bytes"
  4. "fmt"
  5. "net/http"
  6. "net/http/httptest"
  7. "os"
  8. "path/filepath"
  9. "sync/atomic"
  10. "testing"
  11. "time"
  12. "github.com/QuantumNous/new-api/common"
  13. "github.com/gin-gonic/gin"
  14. "github.com/stretchr/testify/assert"
  15. "github.com/stretchr/testify/require"
  16. )
  17. func setupCaptureTestRouter() *gin.Engine {
  18. gin.SetMode(gin.TestMode)
  19. return gin.New()
  20. }
  21. // captureTestCaptureDir 集成测试共享的捕获文件目录
  22. const captureTestDir = "./data/relay-capture"
  23. // --- 非流式请求测试 ---
  24. func TestRelayCapture_NonStreaming_EnabledUser(t *testing.T) {
  25. clearCaptureCache()
  26. SetCaptureEnabled(42, true)
  27. defer SetCaptureEnabled(42, false)
  28. router := setupCaptureTestRouter()
  29. router.Use(func(c *gin.Context) {
  30. c.Set("id", 42) // must be int, not int64 — matches production TokenAuth behavior
  31. c.Set(common.RequestIdKey, "test-non-stream-"+time.Now().Format("150405"))
  32. c.Next()
  33. })
  34. router.Use(RelayCaptureMiddleware())
  35. router.POST("/v1/chat/completions", func(c *gin.Context) {
  36. c.JSON(200, gin.H{
  37. "id": "chatcmpl-test",
  38. "object": "chat.completion",
  39. "choices": []gin.H{{"message": gin.H{"role": "assistant", "content": "Hello!"}}},
  40. })
  41. })
  42. body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}`
  43. req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body)))
  44. req.Header.Set("Authorization", "Bearer sk-test")
  45. req.Header.Set("Content-Type", "application/json")
  46. rec := httptest.NewRecorder()
  47. router.ServeHTTP(rec, req)
  48. assert.Equal(t, 200, rec.Code)
  49. today := time.Now().Format("2006-01-02")
  50. dateDir := filepath.Join(captureTestDir, today)
  51. files, err := os.ReadDir(dateDir)
  52. require.NoError(t, err)
  53. require.True(t, len(files) > 0, "should have capture log files")
  54. logPath := filepath.Join(dateDir, files[len(files)-1].Name())
  55. content, err := os.ReadFile(logPath)
  56. require.NoError(t, err)
  57. contentStr := string(content)
  58. t.Logf("Capture file content:\n%s", contentStr)
  59. assert.Contains(t, contentStr, "=== REQUEST")
  60. assert.Contains(t, contentStr, "POST /v1/chat/completions")
  61. assert.Contains(t, contentStr, "user_id: 42")
  62. assert.Contains(t, contentStr, "Authorization: Bearer sk-test")
  63. assert.Contains(t, contentStr, "=== RESPONSE")
  64. assert.Contains(t, contentStr, "chatcmpl-test")
  65. assert.Contains(t, contentStr, "=== END")
  66. assert.Contains(t, contentStr, "duration_ms=")
  67. }
  68. // --- 流式 SSE 请求测试 ---
  69. func TestRelayCapture_Streaming_EnabledUser(t *testing.T) {
  70. clearCaptureCache()
  71. t.Cleanup(func() { os.RemoveAll(captureTestDir) })
  72. SetCaptureEnabled(42, true)
  73. defer SetCaptureEnabled(42, false)
  74. router := setupCaptureTestRouter()
  75. router.Use(func(c *gin.Context) {
  76. c.Set("id", 42)
  77. c.Set(common.RequestIdKey, "test-stream-"+time.Now().Format("150405"))
  78. c.Next()
  79. })
  80. router.Use(RelayCaptureMiddleware())
  81. router.POST("/v1/chat/completions", func(c *gin.Context) {
  82. c.Header("Content-Type", "text/event-stream")
  83. c.Header("Cache-Control", "no-cache")
  84. c.Header("Connection", "keep-alive")
  85. flusher := c.Writer.(http.Flusher)
  86. chunks := []string{
  87. "data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n",
  88. "data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"world!\"}}]}\n\n",
  89. "data: [DONE]\n\n",
  90. }
  91. for _, chunk := range chunks {
  92. c.Writer.Write([]byte(chunk))
  93. flusher.Flush()
  94. time.Sleep(5 * time.Millisecond)
  95. }
  96. })
  97. body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}],"stream":true}`
  98. req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body)))
  99. req.Header.Set("Authorization", "Bearer sk-test")
  100. req.Header.Set("Content-Type", "application/json")
  101. rec := httptest.NewRecorder()
  102. router.ServeHTTP(rec, req)
  103. assert.Equal(t, 200, rec.Code)
  104. today := time.Now().Format("2006-01-02")
  105. dateDir := filepath.Join(captureTestDir, today)
  106. files, err := os.ReadDir(dateDir)
  107. require.NoError(t, err)
  108. require.True(t, len(files) > 0, "should have capture log files")
  109. logPath := filepath.Join(dateDir, files[len(files)-1].Name())
  110. content, err := os.ReadFile(logPath)
  111. require.NoError(t, err)
  112. contentStr := string(content)
  113. t.Logf("Stream capture file content:\n%s", contentStr)
  114. assert.Contains(t, contentStr, "=== REQUEST")
  115. assert.Contains(t, contentStr, "POST /v1/chat/completions")
  116. assert.Contains(t, contentStr, "user_id: 42")
  117. assert.Contains(t, contentStr, "=== RESPONSE")
  118. assert.Contains(t, contentStr, "Hello ")
  119. assert.Contains(t, contentStr, "world!")
  120. assert.Contains(t, contentStr, "[DONE]")
  121. assert.Contains(t, contentStr, "=== END")
  122. assert.Contains(t, contentStr, "duration_ms=")
  123. }
  124. // --- 未开启抓包用户测试 ---
  125. func TestRelayCapture_DisabledUser(t *testing.T) {
  126. clearCaptureCache()
  127. router := setupCaptureTestRouter()
  128. router.Use(func(c *gin.Context) {
  129. c.Set("id", 99)
  130. c.Set(common.RequestIdKey, "test-disabled-"+time.Now().Format("150405"))
  131. c.Next()
  132. })
  133. router.Use(RelayCaptureMiddleware())
  134. router.POST("/v1/chat/completions", func(c *gin.Context) {
  135. c.JSON(200, gin.H{"result": "ok"})
  136. })
  137. req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{}`)))
  138. req.Header.Set("Authorization", "Bearer sk-test")
  139. rec := httptest.NewRecorder()
  140. router.ServeHTTP(rec, req)
  141. assert.Equal(t, 200, rec.Code)
  142. today := time.Now().Format("2006-01-02")
  143. dateDir := filepath.Join(captureTestDir, today)
  144. if _, err := os.Stat(dateDir); err == nil {
  145. files, _ := os.ReadDir(dateDir)
  146. for _, f := range files {
  147. assert.NotContains(t, f.Name(), "test-disabled", "disabled user should not generate capture files")
  148. }
  149. }
  150. }
  151. // --- 大量数据流测试 ---
  152. func TestRelayCapture_LargeStream(t *testing.T) {
  153. clearCaptureCache()
  154. t.Cleanup(func() { os.RemoveAll(captureTestDir) })
  155. SetCaptureEnabled(42, true)
  156. defer SetCaptureEnabled(42, false)
  157. router := setupCaptureTestRouter()
  158. router.Use(func(c *gin.Context) {
  159. c.Set("id", 42)
  160. c.Set(common.RequestIdKey, "test-large-"+time.Now().Format("150405"))
  161. c.Next()
  162. })
  163. router.Use(RelayCaptureMiddleware())
  164. router.POST("/v1/chat/completions", func(c *gin.Context) {
  165. c.Header("Content-Type", "text/event-stream")
  166. flusher := c.Writer.(http.Flusher)
  167. for i := 0; i < 50; i++ {
  168. chunk := fmt.Sprintf("data: {\"id\":\"%d\",\"choices\":[{\"delta\":{\"content\":\"chunk%d \"}}]}\n\n", i, i)
  169. c.Writer.Write([]byte(chunk))
  170. flusher.Flush()
  171. }
  172. c.Writer.Write([]byte("data: [DONE]\n\n"))
  173. flusher.Flush()
  174. })
  175. req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4","stream":true}`)))
  176. req.Header.Set("Authorization", "Bearer sk-test")
  177. rec := httptest.NewRecorder()
  178. router.ServeHTTP(rec, req)
  179. assert.Equal(t, 200, rec.Code)
  180. today := time.Now().Format("2006-01-02")
  181. dateDir := filepath.Join(captureTestDir, today)
  182. files, err := os.ReadDir(dateDir)
  183. require.NoError(t, err)
  184. require.True(t, len(files) > 0)
  185. logPath := filepath.Join(dateDir, files[len(files)-1].Name())
  186. content, err := os.ReadFile(logPath)
  187. require.NoError(t, err)
  188. contentStr := string(content)
  189. assert.Contains(t, contentStr, "chunk0")
  190. assert.Contains(t, contentStr, "chunk49")
  191. assert.Contains(t, contentStr, "[DONE]")
  192. }
  193. // --- GET 请求测试 ---
  194. func TestRelayCapture_GetRequest(t *testing.T) {
  195. clearCaptureCache()
  196. t.Cleanup(func() { os.RemoveAll(captureTestDir) })
  197. SetCaptureEnabled(42, true)
  198. defer SetCaptureEnabled(42, false)
  199. router := setupCaptureTestRouter()
  200. router.Use(func(c *gin.Context) {
  201. c.Set("id", 42)
  202. c.Set(common.RequestIdKey, "test-get-"+time.Now().Format("150405"))
  203. c.Next()
  204. })
  205. router.Use(RelayCaptureMiddleware())
  206. router.GET("/v1/models", func(c *gin.Context) {
  207. c.JSON(200, gin.H{
  208. "data": []gin.H{{"id": "gpt-4", "object": "model"}},
  209. })
  210. })
  211. req, _ := http.NewRequest("GET", "/v1/models", nil)
  212. req.Header.Set("Authorization", "Bearer sk-test")
  213. rec := httptest.NewRecorder()
  214. router.ServeHTTP(rec, req)
  215. assert.Equal(t, 200, rec.Code)
  216. today := time.Now().Format("2006-01-02")
  217. dateDir := filepath.Join(captureTestDir, today)
  218. files, err := os.ReadDir(dateDir)
  219. require.NoError(t, err)
  220. require.True(t, len(files) > 0)
  221. logPath := filepath.Join(dateDir, files[len(files)-1].Name())
  222. content, err := os.ReadFile(logPath)
  223. require.NoError(t, err)
  224. contentStr := string(content)
  225. assert.Contains(t, contentStr, "=== REQUEST")
  226. assert.Contains(t, contentStr, "GET /v1/models")
  227. assert.Contains(t, contentStr, "=== RESPONSE")
  228. assert.Contains(t, contentStr, "gpt-4")
  229. assert.Contains(t, contentStr, "=== END")
  230. }
  231. // --- 零 ID 用户测试 ---
  232. func TestRelayCapture_ZeroUserID(t *testing.T) {
  233. clearCaptureCache()
  234. router := setupCaptureTestRouter()
  235. router.Use(func(c *gin.Context) {
  236. c.Set(common.RequestIdKey, "test-zero-"+time.Now().Format("150405"))
  237. c.Next()
  238. })
  239. router.Use(RelayCaptureMiddleware())
  240. router.GET("/v1/models", func(c *gin.Context) {
  241. c.JSON(200, gin.H{"data": []gin.H{}})
  242. })
  243. req, _ := http.NewRequest("GET", "/v1/models", nil)
  244. rec := httptest.NewRecorder()
  245. router.ServeHTTP(rec, req)
  246. assert.Equal(t, 200, rec.Code)
  247. }
  248. // --- 验证响应不被篡改 ---
  249. func TestRelayCapture_ResponseUnchanged(t *testing.T) {
  250. clearCaptureCache()
  251. t.Cleanup(func() { os.RemoveAll(captureTestDir) })
  252. SetCaptureEnabled(42, true)
  253. defer SetCaptureEnabled(42, false)
  254. router := setupCaptureTestRouter()
  255. router.Use(func(c *gin.Context) {
  256. c.Set("id", 42)
  257. c.Set(common.RequestIdKey, "test-unchanged-"+time.Now().Format("150405"))
  258. c.Next()
  259. })
  260. router.Use(RelayCaptureMiddleware())
  261. router.POST("/v1/chat/completions", func(c *gin.Context) {
  262. c.JSON(200, gin.H{
  263. "id": "chatcmpl-123",
  264. "choices": []gin.H{{"message": gin.H{"content": "exact response"}}},
  265. })
  266. })
  267. req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4"}`)))
  268. req.Header.Set("Content-Type", "application/json")
  269. rec := httptest.NewRecorder()
  270. router.ServeHTTP(rec, req)
  271. assert.Equal(t, 200, rec.Code)
  272. assert.Contains(t, rec.Body.String(), "exact response")
  273. assert.Contains(t, rec.Body.String(), "chatcmpl-123")
  274. }
  275. // --- 并发安全测试 ---
  276. func TestRelayCapture_ConcurrentRequests(t *testing.T) {
  277. clearCaptureCache()
  278. t.Cleanup(func() { os.RemoveAll(captureTestDir) })
  279. SetCaptureEnabled(42, true)
  280. defer SetCaptureEnabled(42, false)
  281. router := setupCaptureTestRouter()
  282. var reqCounter int64 = 0
  283. router.Use(func(c *gin.Context) {
  284. c.Set("id", 42)
  285. n := atomic.AddInt64(&reqCounter, 1)
  286. c.Set(common.RequestIdKey, fmt.Sprintf("concurrent-%d-%d", n, time.Now().UnixNano()))
  287. c.Next()
  288. })
  289. router.Use(RelayCaptureMiddleware())
  290. router.POST("/v1/chat/completions", func(c *gin.Context) {
  291. c.JSON(200, gin.H{"status": "ok"})
  292. })
  293. done := make(chan bool, 10)
  294. for i := 0; i < 10; i++ {
  295. go func(idx int) {
  296. req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(fmt.Sprintf(`{"req":%d}`, idx))))
  297. req.Header.Set("Content-Type", "application/json")
  298. rec := httptest.NewRecorder()
  299. router.ServeHTTP(rec, req)
  300. assert.Equal(t, 200, rec.Code)
  301. done <- true
  302. }(i)
  303. }
  304. for i := 0; i < 10; i++ {
  305. select {
  306. case <-done:
  307. case <-time.After(10 * time.Second):
  308. t.Fatal("timeout waiting for concurrent request")
  309. }
  310. }
  311. time.Sleep(100 * time.Millisecond)
  312. today := time.Now().Format("2006-01-02")
  313. dateDir := filepath.Join(captureTestDir, today)
  314. files, err := os.ReadDir(dateDir)
  315. require.NoError(t, err)
  316. assert.Equal(t, 10, len(files), "should have 10 capture log files")
  317. for _, f := range files {
  318. content, err := os.ReadFile(filepath.Join(dateDir, f.Name()))
  319. require.NoError(t, err)
  320. contentStr := string(content)
  321. assert.Contains(t, contentStr, "=== REQUEST")
  322. assert.Contains(t, contentStr, "=== RESPONSE")
  323. assert.Contains(t, contentStr, "=== END")
  324. }
  325. }