You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

195 lines
4.4 KiB

  1. package middleware
  2. import (
  3. "fmt"
  4. "io"
  5. "net/http/httptest"
  6. "os"
  7. "path/filepath"
  8. "testing"
  9. "time"
  10. "github.com/gin-gonic/gin"
  11. "github.com/stretchr/testify/assert"
  12. "github.com/stretchr/testify/require"
  13. )
  14. func clearCaptureCache() {
  15. captureEnabledUsers.Range(func(key, value interface{}) bool {
  16. captureEnabledUsers.Delete(key)
  17. return true
  18. })
  19. }
  20. func TestCaptureEnabled_NotFound(t *testing.T) {
  21. clearCaptureCache()
  22. assert.False(t, isCaptureEnabled(1))
  23. }
  24. func TestSetCaptureEnabled_Enable(t *testing.T) {
  25. clearCaptureCache()
  26. SetCaptureEnabled(1, true)
  27. assert.True(t, isCaptureEnabled(1))
  28. SetCaptureEnabled(1, false)
  29. }
  30. func TestSetCaptureEnabled_Disable(t *testing.T) {
  31. clearCaptureCache()
  32. SetCaptureEnabled(1, true)
  33. SetCaptureEnabled(1, false)
  34. assert.False(t, isCaptureEnabled(1))
  35. }
  36. func TestSetCaptureEnabled_DifferentUser(t *testing.T) {
  37. clearCaptureCache()
  38. SetCaptureEnabled(1, true)
  39. assert.False(t, isCaptureEnabled(2))
  40. SetCaptureEnabled(1, false)
  41. }
  42. func TestGetCaptureEnabledUsersCount(t *testing.T) {
  43. clearCaptureCache()
  44. assert.Equal(t, 0, GetCaptureEnabledUsersCount())
  45. SetCaptureEnabled(1, true)
  46. SetCaptureEnabled(2, true)
  47. assert.Equal(t, 2, GetCaptureEnabledUsersCount())
  48. SetCaptureEnabled(1, false)
  49. assert.Equal(t, 1, GetCaptureEnabledUsersCount())
  50. SetCaptureEnabled(2, false)
  51. }
  52. // --- captureResponseWriter 测试 ---
  53. func TestCaptureResponseWriter_Write(t *testing.T) {
  54. gin.SetMode(gin.TestMode)
  55. rec := httptest.NewRecorder()
  56. c, _ := gin.CreateTestContext(rec)
  57. ch := make(chan []byte, 16)
  58. cw := &captureResponseWriter{
  59. ResponseWriter: c.Writer,
  60. ch: ch,
  61. }
  62. n, err := cw.Write([]byte("hello"))
  63. assert.NoError(t, err)
  64. assert.Equal(t, 5, n)
  65. assert.Equal(t, "hello", rec.Body.String())
  66. select {
  67. case data := <-ch:
  68. assert.Equal(t, []byte("hello"), data)
  69. default:
  70. t.Fatal("expected data in channel")
  71. }
  72. }
  73. func TestCaptureResponseWriter_WriteString(t *testing.T) {
  74. gin.SetMode(gin.TestMode)
  75. rec := httptest.NewRecorder()
  76. c, _ := gin.CreateTestContext(rec)
  77. ch := make(chan []byte, 16)
  78. cw := &captureResponseWriter{
  79. ResponseWriter: c.Writer,
  80. ch: ch,
  81. }
  82. n, err := cw.WriteString("test-string")
  83. assert.NoError(t, err)
  84. assert.Equal(t, len("test-string"), n)
  85. assert.Equal(t, "test-string", rec.Body.String())
  86. select {
  87. case data := <-ch:
  88. assert.Equal(t, []byte("test-string"), data)
  89. default:
  90. t.Fatal("expected data in channel")
  91. }
  92. }
  93. // --- startCaptureWriter 测试 ---
  94. func TestStartCaptureWriter_WritesAllChunks(t *testing.T) {
  95. f, err := os.CreateTemp("", "test-capture-*.log")
  96. require.NoError(t, err)
  97. filePath := f.Name()
  98. defer os.Remove(filePath)
  99. start := time.Now()
  100. chIn, doneCh := startCaptureWriter(f, "test-req", start)
  101. // 发送 130 个 chunk(大于 channel buffer 128)
  102. for i := 0; i < 130; i++ {
  103. chIn <- []byte(fmt.Sprintf("data: %d\n", i))
  104. }
  105. close(chIn)
  106. select {
  107. case <-doneCh:
  108. case <-time.After(5 * time.Second):
  109. t.Fatal("timeout waiting for doneCh")
  110. }
  111. content, err := os.ReadFile(filePath)
  112. require.NoError(t, err)
  113. contentStr := string(content)
  114. assert.Contains(t, contentStr, "data: 0")
  115. assert.Contains(t, contentStr, "data: 19")
  116. assert.Contains(t, contentStr, "=== END")
  117. assert.Contains(t, contentStr, "duration_ms=")
  118. }
  119. // --- 辅助函数测试 ---
  120. func TestCreateCaptureFile(t *testing.T) {
  121. f, err := createCaptureFile("test-req-123")
  122. require.NoError(t, err)
  123. defer os.RemoveAll(filepath.Dir(f.Name()))
  124. assert.Contains(t, f.Name(), "relay-capture")
  125. assert.Contains(t, f.Name(), "test-req-123.log")
  126. dir := filepath.Dir(f.Name())
  127. info, err := os.Stat(dir)
  128. require.NoError(t, err)
  129. assert.True(t, info.IsDir())
  130. f.Close()
  131. }
  132. func TestWriteRequestBlock(t *testing.T) {
  133. gin.SetMode(gin.TestMode)
  134. w := httptest.NewRecorder()
  135. c, _ := gin.CreateTestContext(w)
  136. c.Set("id", int64(42))
  137. c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
  138. c.Request.Header.Set("Authorization", "Bearer sk-test")
  139. c.Request.Header.Set("Content-Type", "application/json")
  140. f, err := os.CreateTemp("", "test-req-block-*.log")
  141. require.NoError(t, err)
  142. filePath := f.Name()
  143. defer os.Remove(filePath)
  144. writeRequestBlock(c, f, 42)
  145. f.Seek(0, io.SeekStart)
  146. content, err := io.ReadAll(f)
  147. require.NoError(t, err)
  148. contentStr := string(content)
  149. assert.Contains(t, contentStr, "=== REQUEST")
  150. assert.Contains(t, contentStr, "POST /v1/chat/completions")
  151. assert.Contains(t, contentStr, "user_id: 42")
  152. }