rate_limiter_integration_test.go 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158
  1. //go:build integration
  2. package middleware
  3. import (
  4. "context"
  5. "fmt"
  6. "net/http"
  7. "net/http/httptest"
  8. "os"
  9. "path/filepath"
  10. "strconv"
  11. "testing"
  12. "time"
  13. "github.com/gin-gonic/gin"
  14. "github.com/redis/go-redis/v9"
  15. "github.com/stretchr/testify/require"
  16. tcredis "github.com/testcontainers/testcontainers-go/modules/redis"
  17. )
  18. const redisImageTag = "redis:8.4-alpine"
  19. func TestRateLimiterSetsTTLAndDoesNotRefresh(t *testing.T) {
  20. gin.SetMode(gin.TestMode)
  21. ctx := context.Background()
  22. rdb := startRedis(t, ctx)
  23. limiter := NewRateLimiter(rdb)
  24. router := gin.New()
  25. router.Use(limiter.Limit("ttl-test", 10, 2*time.Second))
  26. router.GET("/test", func(c *gin.Context) {
  27. c.JSON(http.StatusOK, gin.H{"ok": true})
  28. })
  29. recorder := performRequest(router)
  30. require.Equal(t, http.StatusOK, recorder.Code)
  31. redisKey := limiter.prefix + "ttl-test:127.0.0.1"
  32. ttlBefore, err := rdb.PTTL(ctx, redisKey).Result()
  33. require.NoError(t, err)
  34. require.Greater(t, ttlBefore, time.Duration(0))
  35. require.LessOrEqual(t, ttlBefore, 2*time.Second)
  36. time.Sleep(50 * time.Millisecond)
  37. recorder = performRequest(router)
  38. require.Equal(t, http.StatusOK, recorder.Code)
  39. ttlAfter, err := rdb.PTTL(ctx, redisKey).Result()
  40. require.NoError(t, err)
  41. require.Less(t, ttlAfter, ttlBefore)
  42. }
  43. func TestRateLimiterFixesMissingTTL(t *testing.T) {
  44. gin.SetMode(gin.TestMode)
  45. ctx := context.Background()
  46. rdb := startRedis(t, ctx)
  47. limiter := NewRateLimiter(rdb)
  48. router := gin.New()
  49. router.Use(limiter.Limit("ttl-missing", 10, 2*time.Second))
  50. router.GET("/test", func(c *gin.Context) {
  51. c.JSON(http.StatusOK, gin.H{"ok": true})
  52. })
  53. redisKey := limiter.prefix + "ttl-missing:127.0.0.1"
  54. require.NoError(t, rdb.Set(ctx, redisKey, 5, 0).Err())
  55. ttlBefore, err := rdb.PTTL(ctx, redisKey).Result()
  56. require.NoError(t, err)
  57. require.Less(t, ttlBefore, time.Duration(0))
  58. recorder := performRequest(router)
  59. require.Equal(t, http.StatusOK, recorder.Code)
  60. ttlAfter, err := rdb.PTTL(ctx, redisKey).Result()
  61. require.NoError(t, err)
  62. require.Greater(t, ttlAfter, time.Duration(0))
  63. }
  64. func performRequest(router *gin.Engine) *httptest.ResponseRecorder {
  65. req := httptest.NewRequest(http.MethodGet, "/test", nil)
  66. req.RemoteAddr = "127.0.0.1:1234"
  67. recorder := httptest.NewRecorder()
  68. router.ServeHTTP(recorder, req)
  69. return recorder
  70. }
  71. func startRedis(t *testing.T, ctx context.Context) *redis.Client {
  72. t.Helper()
  73. ensureDockerAvailable(t)
  74. redisContainer, err := tcredis.Run(ctx, redisImageTag)
  75. require.NoError(t, err)
  76. t.Cleanup(func() {
  77. _ = redisContainer.Terminate(ctx)
  78. })
  79. redisHost, err := redisContainer.Host(ctx)
  80. require.NoError(t, err)
  81. redisPort, err := redisContainer.MappedPort(ctx, "6379/tcp")
  82. require.NoError(t, err)
  83. rdb := redis.NewClient(&redis.Options{
  84. Addr: fmt.Sprintf("%s:%d", redisHost, redisPort.Int()),
  85. DB: 0,
  86. })
  87. require.NoError(t, rdb.Ping(ctx).Err())
  88. t.Cleanup(func() {
  89. _ = rdb.Close()
  90. })
  91. return rdb
  92. }
  93. func ensureDockerAvailable(t *testing.T) {
  94. t.Helper()
  95. if dockerAvailable() {
  96. return
  97. }
  98. t.Skip("Docker 未启用,跳过依赖 testcontainers 的集成测试")
  99. }
  100. func dockerAvailable() bool {
  101. if os.Getenv("DOCKER_HOST") != "" {
  102. return true
  103. }
  104. socketCandidates := []string{
  105. "/var/run/docker.sock",
  106. filepath.Join(os.Getenv("XDG_RUNTIME_DIR"), "docker.sock"),
  107. filepath.Join(userHomeDir(), ".docker", "run", "docker.sock"),
  108. filepath.Join(userHomeDir(), ".docker", "desktop", "docker.sock"),
  109. filepath.Join("/run/user", strconv.Itoa(os.Getuid()), "docker.sock"),
  110. }
  111. for _, socket := range socketCandidates {
  112. if socket == "" {
  113. continue
  114. }
  115. if _, err := os.Stat(socket); err == nil {
  116. return true
  117. }
  118. }
  119. return false
  120. }
  121. func userHomeDir() string {
  122. home, err := os.UserHomeDir()
  123. if err != nil {
  124. return ""
  125. }
  126. return home
  127. }