client_request_id_test.go 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. package middleware
  2. import (
  3. "context"
  4. "net/http"
  5. "net/http/httptest"
  6. "strings"
  7. "testing"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  9. "github.com/gin-gonic/gin"
  10. "github.com/stretchr/testify/require"
  11. )
  12. func TestClientRequestIDGeneratesAndExposesID(t *testing.T) {
  13. gin.SetMode(gin.TestMode)
  14. router := gin.New()
  15. router.Use(ClientRequestID())
  16. router.GET("/", func(c *gin.Context) {
  17. value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
  18. c.String(http.StatusOK, value)
  19. })
  20. req := httptest.NewRequest(http.MethodGet, "/", nil)
  21. w := httptest.NewRecorder()
  22. router.ServeHTTP(w, req)
  23. require.Equal(t, http.StatusOK, w.Code)
  24. require.NotEmpty(t, w.Body.String())
  25. require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader))
  26. }
  27. func TestClientRequestIDBoundsExistingContextID(t *testing.T) {
  28. gin.SetMode(gin.TestMode)
  29. router := gin.New()
  30. router.Use(ClientRequestID())
  31. router.GET("/", func(c *gin.Context) {
  32. value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
  33. c.String(http.StatusOK, value)
  34. })
  35. req := httptest.NewRequest(http.MethodGet, "/", nil)
  36. req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, strings.Repeat("x", 200)))
  37. w := httptest.NewRecorder()
  38. router.ServeHTTP(w, req)
  39. require.Len(t, w.Body.String(), 36)
  40. require.NotEqual(t, strings.Repeat("x", maxPersistentRequestIDBytes), w.Body.String())
  41. require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader))
  42. }
  43. func TestClientRequestIDPreservesExistingContextID(t *testing.T) {
  44. gin.SetMode(gin.TestMode)
  45. router := gin.New()
  46. router.Use(ClientRequestID())
  47. router.GET("/", func(c *gin.Context) {
  48. value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
  49. c.String(http.StatusOK, value)
  50. })
  51. w := httptest.NewRecorder()
  52. req := httptest.NewRequest(http.MethodGet, "/", nil)
  53. req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "existing-client-request-id"))
  54. router.ServeHTTP(w, req)
  55. require.Equal(t, http.StatusOK, w.Code)
  56. require.Equal(t, "existing-client-request-id", w.Body.String())
  57. require.Equal(t, "existing-client-request-id", w.Header().Get(clientRequestIDHeader))
  58. }