new-api/middleware/relay_capture_test.go

package middleware

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

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

func TestCaptureResponseWriter_Write(t *testing.T) {
	gin.SetMode(gin.TestMode)
	rec := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(rec)
	collector := newResponseCaptureCollector()

	cw := &captureResponseWriter{
		ResponseWriter: c.Writer,
		collector:      collector,
	}

	n, err := cw.Write([]byte("hello"))
	assert.NoError(t, err)
	assert.Equal(t, 5, n)
	assert.Equal(t, "hello", rec.Body.String())

	assert.Equal(t, "hello", collector.Body())
	assert.Equal(t, []string{"hello"}, collector.Chunks())
}

func TestCaptureResponseWriter_WriteString(t *testing.T) {
	gin.SetMode(gin.TestMode)
	rec := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(rec)
	collector := newResponseCaptureCollector()

	cw := &captureResponseWriter{
		ResponseWriter: c.Writer,
		collector:      collector,
	}

	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())

	assert.Equal(t, "test-string", collector.Body())
	assert.Equal(t, []string{"test-string"}, collector.Chunks())
}

type partialResponseWriter struct {
	gin.ResponseWriter
}

func (w partialResponseWriter) Write(b []byte) (int, error) {
	_, err := w.ResponseWriter.Write(b[:2])
	return 2, err
}

func TestCaptureResponseWriter_RecordsOnlyWrittenBytes(t *testing.T) {
	gin.SetMode(gin.TestMode)
	rec := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(rec)
	collector := newResponseCaptureCollector()

	cw := &captureResponseWriter{
		ResponseWriter: partialResponseWriter{ResponseWriter: c.Writer},
		collector:      collector,
	}

	n, err := cw.Write([]byte("hello"))
	assert.NoError(t, err)
	assert.Equal(t, 2, n)
	assert.Equal(t, "he", rec.Body.String())
	assert.Equal(t, "he", collector.Body())
	assert.Equal(t, []string{"he"}, collector.Chunks())
}

func TestCreateCaptureFile(t *testing.T) {
	f, err := createCaptureFile("test-req-123")
	require.NoError(t, err)
	defer os.RemoveAll(filepath.Dir(f.Name()))
	defer f.Close()

	assert.Contains(t, f.Name(), "relay-capture")
	assert.Contains(t, f.Name(), "test-req-123.json")
	assert.NotContains(t, f.Name(), ".log")

	dir := filepath.Dir(f.Name())
	info, err := os.Stat(dir)
	require.NoError(t, err)
	assert.True(t, info.IsDir())
}

func TestNormalizeCaptureRequestID_UsesExistingID(t *testing.T) {
	requestID := normalizeCaptureRequestID("req-123")
	assert.Equal(t, "req-123", requestID)
}

func TestNormalizeCaptureRequestID_FallbackForEmptyID(t *testing.T) {
	requestID := normalizeCaptureRequestID("")
	assert.True(t, strings.HasPrefix(requestID, "capture-"))
	assert.False(t, strings.HasSuffix(requestID, ".json"))
}

func TestCopyHeadersForCapture_RedactsSensitiveRequestHeaders(t *testing.T) {
	headers := http.Header{
		"Authorization":       []string{"Bearer sk-test"},
		"authorization":       []string{"Bearer lower"},
		"X-Api-Key":           []string{"x-key"},
		"Api-Key":             []string{"api-key"},
		"Cookie":              []string{"session=secret"},
		"Proxy-Authorization": []string{"Basic secret"},
		"Content-Type":        []string{"application/json"},
		"X-Trace-ID":          []string{"trace-1"},
	}

	copied := copyHeadersForCapture(headers, true)

	assert.Equal(t, []string{"[REDACTED]"}, copied["Authorization"])
	assert.Equal(t, []string{"[REDACTED]"}, copied["authorization"])
	assert.Equal(t, []string{"[REDACTED]"}, copied["X-Api-Key"])
	assert.Equal(t, []string{"[REDACTED]"}, copied["Api-Key"])
	assert.Equal(t, []string{"[REDACTED]"}, copied["Cookie"])
	assert.Equal(t, []string{"[REDACTED]"}, copied["Proxy-Authorization"])
	assert.Equal(t, []string{"application/json"}, copied["Content-Type"])
	assert.Equal(t, []string{"trace-1"}, copied["X-Trace-ID"])

	assert.Equal(t, []string{"Bearer sk-test"}, headers["Authorization"])
}

func TestCopyHeadersForCapture_ResponseHeadersNotRedacted(t *testing.T) {
	headers := http.Header{}
	headers.Add("Authorization", "response-token")
	headers.Add("Content-Type", "application/json")

	copied := copyHeadersForCapture(headers, false)

	assert.Equal(t, []string{"response-token"}, copied["Authorization"])
	assert.Equal(t, []string{"application/json"}, copied["Content-Type"])
}

func TestBuildCaptureRequestRecord_PostBody(t *testing.T) {
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	body := `{"model":"gpt-4"}`
	c.Request = httptest.NewRequest("POST", "/v1/chat/completions?stream=true", bytes.NewReader([]byte(body)))
	c.Request.Header.Set("Authorization", "Bearer sk-test")
	c.Request.Header.Set("Content-Type", "application/json")

	var captureErrors []string
	record := buildCaptureRequestRecord(c, &captureErrors)

	assert.Equal(t, "POST", record.Method)
	assert.Equal(t, "/v1/chat/completions", record.Path)
	assert.Equal(t, "stream=true", record.Query)
	assert.Equal(t, body, record.Body)
	assert.Equal(t, []string{"[REDACTED]"}, record.Headers["Authorization"])
	assert.Equal(t, []string{"application/json"}, record.Headers["Content-Type"])
	assert.Empty(t, captureErrors)
}

func TestBuildCaptureRequestRecord_GetNoBody(t *testing.T) {
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest("GET", "/v1/models", nil)

	var captureErrors []string
	record := buildCaptureRequestRecord(c, &captureErrors)

	assert.Equal(t, "GET", record.Method)
	assert.Equal(t, "/v1/models", record.Path)
	assert.Equal(t, "", record.Query)
	assert.Equal(t, "", record.Body)
	assert.Empty(t, captureErrors)
}