api_key_cache_integration_test.go 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. //go:build integration
  2. package repository
  3. import (
  4. "context"
  5. "fmt"
  6. "testing"
  7. "time"
  8. "github.com/redis/go-redis/v9"
  9. "github.com/stretchr/testify/require"
  10. "github.com/stretchr/testify/suite"
  11. )
  12. type ApiKeyCacheSuite struct {
  13. IntegrationRedisSuite
  14. }
  15. func (s *ApiKeyCacheSuite) TestCreateAttemptCount() {
  16. tests := []struct {
  17. name string
  18. fn func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache)
  19. }{
  20. {
  21. name: "missing_key_returns_zero_nil",
  22. fn: func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache) {
  23. userID := int64(1)
  24. count, err := cache.GetCreateAttemptCount(ctx, userID)
  25. require.NoError(s.T(), err, "expected nil error for missing key")
  26. require.Equal(s.T(), 0, count, "expected zero count for missing key")
  27. },
  28. },
  29. {
  30. name: "increment_increases_count_and_sets_ttl",
  31. fn: func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache) {
  32. userID := int64(1)
  33. key := fmt.Sprintf("%s%d", apiKeyRateLimitKeyPrefix, userID)
  34. require.NoError(s.T(), cache.IncrementCreateAttemptCount(ctx, userID), "IncrementCreateAttemptCount")
  35. require.NoError(s.T(), cache.IncrementCreateAttemptCount(ctx, userID), "IncrementCreateAttemptCount 2")
  36. count, err := cache.GetCreateAttemptCount(ctx, userID)
  37. require.NoError(s.T(), err, "GetCreateAttemptCount")
  38. require.Equal(s.T(), 2, count, "count mismatch")
  39. ttl, err := rdb.TTL(ctx, key).Result()
  40. require.NoError(s.T(), err, "TTL")
  41. s.AssertTTLWithin(ttl, 1*time.Second, apiKeyRateLimitDuration)
  42. },
  43. },
  44. {
  45. name: "delete_removes_key",
  46. fn: func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache) {
  47. userID := int64(1)
  48. require.NoError(s.T(), cache.IncrementCreateAttemptCount(ctx, userID))
  49. require.NoError(s.T(), cache.DeleteCreateAttemptCount(ctx, userID), "DeleteCreateAttemptCount")
  50. count, err := cache.GetCreateAttemptCount(ctx, userID)
  51. require.NoError(s.T(), err, "expected nil error after delete")
  52. require.Equal(s.T(), 0, count, "expected zero count after delete")
  53. },
  54. },
  55. }
  56. for _, tt := range tests {
  57. s.Run(tt.name, func() {
  58. // 每个 case 重新获取隔离资源
  59. rdb := testRedis(s.T())
  60. cache := &apiKeyCache{rdb: rdb}
  61. ctx := context.Background()
  62. tt.fn(ctx, rdb, cache)
  63. })
  64. }
  65. }
  66. func (s *ApiKeyCacheSuite) TestDailyUsage() {
  67. tests := []struct {
  68. name string
  69. fn func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache)
  70. }{
  71. {
  72. name: "increment_increases_count",
  73. fn: func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache) {
  74. dailyKey := "daily:sk-test"
  75. require.NoError(s.T(), cache.IncrementDailyUsage(ctx, dailyKey), "IncrementDailyUsage")
  76. require.NoError(s.T(), cache.IncrementDailyUsage(ctx, dailyKey), "IncrementDailyUsage 2")
  77. n, err := rdb.Get(ctx, dailyKey).Int()
  78. require.NoError(s.T(), err, "Get dailyKey")
  79. require.Equal(s.T(), 2, n, "expected daily usage=2")
  80. },
  81. },
  82. {
  83. name: "set_expiry_sets_ttl",
  84. fn: func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache) {
  85. dailyKey := "daily:sk-test-expiry"
  86. require.NoError(s.T(), cache.IncrementDailyUsage(ctx, dailyKey))
  87. require.NoError(s.T(), cache.SetDailyUsageExpiry(ctx, dailyKey, 1*time.Hour), "SetDailyUsageExpiry")
  88. ttl, err := rdb.TTL(ctx, dailyKey).Result()
  89. require.NoError(s.T(), err, "TTL dailyKey")
  90. require.Greater(s.T(), ttl, time.Duration(0), "expected ttl > 0")
  91. },
  92. },
  93. }
  94. for _, tt := range tests {
  95. s.Run(tt.name, func() {
  96. rdb := testRedis(s.T())
  97. cache := &apiKeyCache{rdb: rdb}
  98. ctx := context.Background()
  99. tt.fn(ctx, rdb, cache)
  100. })
  101. }
  102. }
  103. func TestApiKeyCacheSuite(t *testing.T) {
  104. suite.Run(t, new(ApiKeyCacheSuite))
  105. }