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