misc_coverage_test.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126
  1. //go:build unit
  2. package middleware
  3. import (
  4. "bytes"
  5. "context"
  6. "io"
  7. "net/http"
  8. "net/http/httptest"
  9. "testing"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  11. "github.com/Wei-Shaw/sub2api/internal/service"
  12. "github.com/gin-gonic/gin"
  13. "github.com/stretchr/testify/require"
  14. )
  15. func TestClientRequestID_GeneratesWhenMissing(t *testing.T) {
  16. gin.SetMode(gin.TestMode)
  17. r := gin.New()
  18. r.Use(ClientRequestID())
  19. r.GET("/t", func(c *gin.Context) {
  20. v := c.Request.Context().Value(ctxkey.ClientRequestID)
  21. require.NotNil(t, v)
  22. id, ok := v.(string)
  23. require.True(t, ok)
  24. require.NotEmpty(t, id)
  25. c.Status(http.StatusOK)
  26. })
  27. w := httptest.NewRecorder()
  28. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  29. r.ServeHTTP(w, req)
  30. require.Equal(t, http.StatusOK, w.Code)
  31. }
  32. func TestClientRequestID_PreservesExisting(t *testing.T) {
  33. gin.SetMode(gin.TestMode)
  34. r := gin.New()
  35. r.Use(ClientRequestID())
  36. r.GET("/t", func(c *gin.Context) {
  37. id, ok := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
  38. require.True(t, ok)
  39. require.Equal(t, "keep", id)
  40. c.Status(http.StatusOK)
  41. })
  42. w := httptest.NewRecorder()
  43. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  44. req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, "keep"))
  45. r.ServeHTTP(w, req)
  46. require.Equal(t, http.StatusOK, w.Code)
  47. }
  48. func TestRequestBodyLimit_LimitsBody(t *testing.T) {
  49. gin.SetMode(gin.TestMode)
  50. r := gin.New()
  51. r.Use(RequestBodyLimit(4))
  52. r.POST("/t", func(c *gin.Context) {
  53. _, err := io.ReadAll(c.Request.Body)
  54. require.Error(t, err)
  55. c.Status(http.StatusOK)
  56. })
  57. w := httptest.NewRecorder()
  58. req := httptest.NewRequest(http.MethodPost, "/t", bytes.NewBufferString("12345"))
  59. r.ServeHTTP(w, req)
  60. require.Equal(t, http.StatusOK, w.Code)
  61. }
  62. func TestForcePlatform_SetsContextAndGinValue(t *testing.T) {
  63. gin.SetMode(gin.TestMode)
  64. r := gin.New()
  65. r.Use(ForcePlatform("anthropic"))
  66. r.GET("/t", func(c *gin.Context) {
  67. require.True(t, HasForcePlatform(c))
  68. v, ok := GetForcePlatformFromContext(c)
  69. require.True(t, ok)
  70. require.Equal(t, "anthropic", v)
  71. ctxV := c.Request.Context().Value(ctxkey.ForcePlatform)
  72. require.Equal(t, "anthropic", ctxV)
  73. c.Status(http.StatusOK)
  74. })
  75. w := httptest.NewRecorder()
  76. req := httptest.NewRequest(http.MethodGet, "/t", nil)
  77. r.ServeHTTP(w, req)
  78. require.Equal(t, http.StatusOK, w.Code)
  79. }
  80. func TestAuthSubjectHelpers_RoundTrip(t *testing.T) {
  81. c := &gin.Context{}
  82. c.Set(string(ContextKeyUser), AuthSubject{UserID: 1, Concurrency: 2})
  83. c.Set(string(ContextKeyUserRole), "admin")
  84. sub, ok := GetAuthSubjectFromContext(c)
  85. require.True(t, ok)
  86. require.Equal(t, int64(1), sub.UserID)
  87. require.Equal(t, 2, sub.Concurrency)
  88. role, ok := GetUserRoleFromContext(c)
  89. require.True(t, ok)
  90. require.Equal(t, "admin", role)
  91. }
  92. func TestAPIKeyAndSubscriptionFromContext(t *testing.T) {
  93. c := &gin.Context{}
  94. key := &service.APIKey{ID: 1}
  95. c.Set(string(ContextKeyAPIKey), key)
  96. gotKey, ok := GetAPIKeyFromContext(c)
  97. require.True(t, ok)
  98. require.Equal(t, int64(1), gotKey.ID)
  99. sub := &service.UserSubscription{ID: 2}
  100. c.Set(string(ContextKeySubscription), sub)
  101. gotSub, ok := GetSubscriptionFromContext(c)
  102. require.True(t, ok)
  103. require.Equal(t, int64(2), gotSub.ID)
  104. }