package middleware import ( "fmt" "io" "net/http/httptest" "os" "path/filepath" "testing" "time" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func clearCaptureCache() { captureEnabledUsers.Range(func(key, value interface{}) bool { captureEnabledUsers.Delete(key) return true }) } func TestCaptureEnabled_NotFound(t *testing.T) { clearCaptureCache() assert.False(t, isCaptureEnabled(1)) } func TestSetCaptureEnabled_Enable(t *testing.T) { clearCaptureCache() SetCaptureEnabled(1, true) assert.True(t, isCaptureEnabled(1)) SetCaptureEnabled(1, false) } func TestSetCaptureEnabled_Disable(t *testing.T) { clearCaptureCache() SetCaptureEnabled(1, true) SetCaptureEnabled(1, false) assert.False(t, isCaptureEnabled(1)) } func TestSetCaptureEnabled_DifferentUser(t *testing.T) { clearCaptureCache() SetCaptureEnabled(1, true) assert.False(t, isCaptureEnabled(2)) SetCaptureEnabled(1, false) } func TestGetCaptureEnabledUsersCount(t *testing.T) { clearCaptureCache() assert.Equal(t, 0, GetCaptureEnabledUsersCount()) SetCaptureEnabled(1, true) SetCaptureEnabled(2, true) assert.Equal(t, 2, GetCaptureEnabledUsersCount()) SetCaptureEnabled(1, false) assert.Equal(t, 1, GetCaptureEnabledUsersCount()) SetCaptureEnabled(2, false) } // --- captureResponseWriter 测试 --- func TestCaptureResponseWriter_Write(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) ch := make(chan []byte, 16) cw := &captureResponseWriter{ ResponseWriter: c.Writer, ch: ch, } n, err := cw.Write([]byte("hello")) assert.NoError(t, err) assert.Equal(t, 5, n) assert.Equal(t, "hello", rec.Body.String()) select { case data := <-ch: assert.Equal(t, []byte("hello"), data) default: t.Fatal("expected data in channel") } } func TestCaptureResponseWriter_WriteString(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) ch := make(chan []byte, 16) cw := &captureResponseWriter{ ResponseWriter: c.Writer, ch: ch, } n, err := cw.WriteString("test-string") assert.NoError(t, err) assert.Equal(t, len("test-string"), n) assert.Equal(t, "test-string", rec.Body.String()) select { case data := <-ch: assert.Equal(t, []byte("test-string"), data) default: t.Fatal("expected data in channel") } } // --- startCaptureWriter 测试 --- func TestStartCaptureWriter_WritesAllChunks(t *testing.T) { f, err := os.CreateTemp("", "test-capture-*.log") require.NoError(t, err) filePath := f.Name() defer os.Remove(filePath) start := time.Now() chIn, doneCh := startCaptureWriter(f, "test-req", start) // 发送 130 个 chunk(大于 channel buffer 128) for i := 0; i < 130; i++ { chIn <- []byte(fmt.Sprintf("data: %d\n", i)) } close(chIn) select { case <-doneCh: case <-time.After(5 * time.Second): t.Fatal("timeout waiting for doneCh") } content, err := os.ReadFile(filePath) require.NoError(t, err) contentStr := string(content) assert.Contains(t, contentStr, "data: 0") assert.Contains(t, contentStr, "data: 19") assert.Contains(t, contentStr, "=== END") assert.Contains(t, contentStr, "duration_ms=") } // --- 辅助函数测试 --- func TestCreateCaptureFile(t *testing.T) { f, err := createCaptureFile("test-req-123") require.NoError(t, err) defer os.RemoveAll(filepath.Dir(f.Name())) assert.Contains(t, f.Name(), "relay-capture") assert.Contains(t, f.Name(), "test-req-123.log") dir := filepath.Dir(f.Name()) info, err := os.Stat(dir) require.NoError(t, err) assert.True(t, info.IsDir()) f.Close() } func TestWriteRequestBlock(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Set("id", int64(42)) c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil) c.Request.Header.Set("Authorization", "Bearer sk-test") c.Request.Header.Set("Content-Type", "application/json") f, err := os.CreateTemp("", "test-req-block-*.log") require.NoError(t, err) filePath := f.Name() defer os.Remove(filePath) writeRequestBlock(c, f, 42) f.Seek(0, io.SeekStart) content, err := io.ReadAll(f) require.NoError(t, err) contentStr := string(content) assert.Contains(t, contentStr, "=== REQUEST") assert.Contains(t, contentStr, "POST /v1/chat/completions") assert.Contains(t, contentStr, "user_id: 42") }