// Package middleware provides HTTP middleware components for the CLI Proxy API server. // This file contains unit tests for the middleware components. package middleware import ( "errors" "net/http" "net/http/httptest" "testing" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" ) func setupTestRouter() *gin.Engine { gin.SetMode(gin.TestMode) return gin.New() } func TestCorrelationIDMiddleware(t *testing.T) { router := setupTestRouter() router.Use(CorrelationIDMiddleware()) router.GET("/test", func(c *gin.Context) { id := GetCorrelationID(c) c.JSON(200, gin.H{"correlation_id": id}) }) t.Run("generates correlation ID when not provided", func(t *testing.T) { w := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/test", nil) router.ServeHTTP(w, req) assert.Equal(t, 200, w.Code) // Check that response contains a correlation ID assert.Contains(t, w.Body.String(), "correlation_id") // Check that response header contains the correlation ID assert.NotEmpty(t, w.Header().Get(CorrelationIDHeader)) }) t.Run("uses existing correlation ID from header", func(t *testing.T) { w := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/test", nil) req.Header.Set(CorrelationIDHeader, "test-correlation-id-123") router.ServeHTTP(w, req) assert.Equal(t, 200, w.Code) assert.Contains(t, w.Body.String(), "test-correlation-id-123") assert.Equal(t, "test-correlation-id-123", w.Header().Get(CorrelationIDHeader)) }) } func TestErrorHandlerMiddleware(t *testing.T) { router := setupTestRouter() router.Use(CorrelationIDMiddleware()) router.Use(ErrorHandlerMiddleware()) router.GET("/error", func(c *gin.Context) { c.Error(errors.New("test error")) c.Abort() }) router.GET("/success", func(c *gin.Context) { c.JSON(200, gin.H{"status": "ok"}) }) t.Run("handles errors gracefully", func(t *testing.T) { w := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/error", nil) router.ServeHTTP(w, req) assert.Equal(t, 500, w.Code) assert.Contains(t, w.Body.String(), "INTERNAL_ERROR") }) t.Run("passes through successful requests", func(t *testing.T) { w := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/success", nil) router.ServeHTTP(w, req) assert.Equal(t, 200, w.Code) assert.Contains(t, w.Body.String(), "ok") }) } func TestRecoveryMiddleware(t *testing.T) { router := setupTestRouter() router.Use(RecoveryMiddleware(nil)) router.Use(CorrelationIDMiddleware()) router.Use(ErrorHandlerMiddleware()) router.GET("/panic", func(c *gin.Context) { panic("test panic") }) router.GET("/normal", func(c *gin.Context) { c.JSON(200, gin.H{"status": "ok"}) }) t.Run("recovers from panic", func(t *testing.T) { w := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/panic", nil) router.ServeHTTP(w, req) assert.Equal(t, 500, w.Code) assert.Contains(t, w.Body.String(), "INTERNAL_ERROR") }) t.Run("normal requests work after panic recovery", func(t *testing.T) { w := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/normal", nil) router.ServeHTTP(w, req) assert.Equal(t, 200, w.Code) assert.Contains(t, w.Body.String(), "ok") }) } func TestSafeHandler(t *testing.T) { router := setupTestRouter() router.Use(CorrelationIDMiddleware()) router.Use(ErrorHandlerMiddleware()) router.GET("/safe-panic", SafeHandler(func(c *gin.Context) { panic("safe handler panic") })) t.Run("safe handler recovers from panic", func(t *testing.T) { w := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/safe-panic", nil) router.ServeHTTP(w, req) assert.Equal(t, 500, w.Code) assert.Contains(t, w.Body.String(), "INTERNAL_ERROR") }) }