mirror of
https://github.com/andwati/telecom-platform.git
synced 2026-10-26 14:42:15 +00:00
feat: Add comprehensive Redis rate limiter tests with mock client, middleware validation, and multi-tier support
- Add MockRedisClient struct with Eval/Del/ZCount methods for Redis operation mocking - Add TestRedisRateLimiter_Allow with under/over limit scenarios and mock Redis response validation - Add TestRedisRateLimitMiddleware with IP-based rate limiting and empty key bypass tests - Add TestIPKeyExtractor/TestUserKeyExtractor/TestAPIKeyExtractor with valid/invalid input validation - Add TestEnd
This commit is contained in:
nutcas3
@@ -0,0 +1,339 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
// MockRedisClient is a mock for Redis client
|
||||
type MockRedisClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockRedisClient) Eval(ctx context.Context, script string, keys []string, args ...interface{}) *redis.Cmd {
|
||||
args = append([]interface{}{script, keys}, args...)
|
||||
return m.Called(args...).Get(0).(*redis.Cmd)
|
||||
}
|
||||
|
||||
func (m *MockRedisClient) Del(ctx context.Context, keys ...string) *redis.IntCmd {
|
||||
return m.Called(keys...).Get(0).(*redis.IntCmd)
|
||||
}
|
||||
|
||||
func (m *MockRedisClient) ZCount(ctx context.Context, key, min, max string) *redis.IntCmd {
|
||||
return m.Called(key, min, max).Get(0).(*redis.IntCmd)
|
||||
}
|
||||
|
||||
func TestRedisRateLimiter_Allow(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
limit int64
|
||||
window time.Duration
|
||||
expected bool
|
||||
}{
|
||||
{"Under limit", 5, time.Minute, true},
|
||||
{"Over limit", 1, time.Minute, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Create mock Redis client
|
||||
mockClient := new(MockRedisClient)
|
||||
|
||||
// Mock Redis response based on test case
|
||||
var mockResult []interface{}
|
||||
if tt.expected {
|
||||
mockResult = []interface{}{int64(1), int64(tt.limit - 1), int64(time.Now().Unix() + 60)}
|
||||
} else {
|
||||
mockResult = []interface{}{int64(0), int64(0), int64(time.Now().Unix() + 60), int64(30)}
|
||||
}
|
||||
|
||||
cmd := redis.NewCmd(context.Background())
|
||||
cmd.SetVal(mockResult)
|
||||
mockClient.On("Eval", mock.AnythingOfType("[]interface {}")).Return(cmd)
|
||||
|
||||
// Create rate limiter
|
||||
limiter := NewRedisRateLimiter(&redis.Client{}, "test", tt.limit, tt.window)
|
||||
limiter.client = mockClient // Replace with mock
|
||||
|
||||
// Test rate limiting
|
||||
result, err := limiter.Allow(context.Background(), "test-key")
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, tt.expected, result.Allowed)
|
||||
assert.Equal(t, tt.limit, result.TotalLimit)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisRateLimitMiddleware(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
keyExtractor func(*gin.Context) string
|
||||
expectedStatus int
|
||||
shouldHaveHeaders bool
|
||||
}{
|
||||
{
|
||||
name: "IP-based rate limiting - allowed",
|
||||
keyExtractor: func(c *gin.Context) string {
|
||||
return "127.0.0.1"
|
||||
},
|
||||
expectedStatus: http.StatusOK,
|
||||
shouldHaveHeaders: true,
|
||||
},
|
||||
{
|
||||
name: "Empty key - bypass rate limiting",
|
||||
keyExtractor: func(c *gin.Context) string {
|
||||
return ""
|
||||
},
|
||||
expectedStatus: http.StatusOK,
|
||||
shouldHaveHeaders: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Create mock Redis client
|
||||
mockClient := new(MockRedisClient)
|
||||
|
||||
// Mock successful rate limit check
|
||||
cmd := redis.NewCmd(context.Background())
|
||||
cmd.SetVal([]interface{}{int64(1), int64(4), int64(time.Now().Unix() + 60)})
|
||||
mockClient.On("Eval", mock.AnythingOfType("[]interface {}")).Return(cmd)
|
||||
|
||||
// Create rate limiter
|
||||
limiter := NewRedisRateLimiter(&redis.Client{}, "test", 5, time.Minute)
|
||||
limiter.client = mockClient // Replace with mock
|
||||
|
||||
// Create Gin router with middleware
|
||||
router := gin.New()
|
||||
router.Use(RedisRateLimitMiddleware(limiter, tt.keyExtractor))
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"message": "success"})
|
||||
})
|
||||
|
||||
// Create request
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
req.RemoteAddr = "127.0.0.1:12345"
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
// Perform request
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
// Check response
|
||||
assert.Equal(t, tt.expectedStatus, w.Code)
|
||||
|
||||
// Check rate limit headers
|
||||
if tt.shouldHaveHeaders {
|
||||
assert.Contains(t, w.Header(), "X-RateLimit-Limit")
|
||||
assert.Contains(t, w.Header(), "X-RateLimit-Remaining")
|
||||
assert.Contains(t, w.Header(), "X-RateLimit-Reset")
|
||||
assert.Equal(t, "5", w.Header().Get("X-RateLimit-Limit"))
|
||||
} else {
|
||||
assert.NotContains(t, w.Header(), "X-RateLimit-Limit")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIPKeyExtractor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
clientIP string
|
||||
expected string
|
||||
}{
|
||||
{"Valid IP", "192.168.1.1", "192.168.1.1"},
|
||||
{"IPv6", "::1", "::1"},
|
||||
{"Empty IP", "", ""},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.RemoteAddr = tt.clientIP + ":12345"
|
||||
|
||||
result := IPKeyExtractor(c)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserKeyExtractor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
userID interface{}
|
||||
clientIP string
|
||||
expected string
|
||||
}{
|
||||
{"With user ID", "user123", "192.168.1.1", "user:user123"},
|
||||
{"Without user ID", nil, "192.168.1.1", "192.168.1.1"},
|
||||
{"Invalid user ID type", 123, "192.168.1.1", "192.168.1.1"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.RemoteAddr = tt.clientIP + ":12345"
|
||||
|
||||
if tt.userID != nil {
|
||||
c.Set("user_id", tt.userID)
|
||||
}
|
||||
|
||||
result := UserKeyExtractor(c)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIKeyExtractor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
apiKey string
|
||||
userID string
|
||||
clientIP string
|
||||
expected string
|
||||
}{
|
||||
{"With API key", "api-key-123", "user123", "192.168.1.1", "api:api-key-123"},
|
||||
{"Without API key", "", "user123", "192.168.1.1", "user:user123"},
|
||||
{"No API key or user", "", "", "192.168.1.1", "192.168.1.1"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.RemoteAddr = tt.clientIP + ":12345"
|
||||
|
||||
if tt.apiKey != "" {
|
||||
c.Request.Header.Set("X-API-Key", tt.apiKey)
|
||||
}
|
||||
if tt.userID != "" {
|
||||
c.Set("user_id", tt.userID)
|
||||
}
|
||||
|
||||
result := APIKeyExtractor(c)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEndpointKeyExtractor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
baseExtractor := func(c *gin.Context) string {
|
||||
return "user123"
|
||||
}
|
||||
|
||||
extractor := EndpointKeyExtractor(baseExtractor)
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/api/v1/users", nil)
|
||||
c.Request.RemoteAddr = "192.168.1.1:12345"
|
||||
|
||||
result := extractor(c)
|
||||
expected := "user123:GET:/api/v1/users"
|
||||
assert.Equal(t, expected, result)
|
||||
}
|
||||
|
||||
func TestMultiTierRateLimiter(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
// Create mock Redis client
|
||||
mockClient := new(MockRedisClient)
|
||||
|
||||
// Mock different tier responses
|
||||
freeCmd := redis.NewCmd(context.Background())
|
||||
freeCmd.SetVal([]interface{}{int64(1), int64(9), int64(time.Now().Unix() + 60)})
|
||||
|
||||
premiumCmd := redis.NewCmd(context.Background())
|
||||
premiumCmd.SetVal([]interface{}{int64(1), int64(99), int64(time.Now().Unix() + 60)})
|
||||
|
||||
mockClient.On("Eval", mock.AnythingOfType("[]interface {}")).Return(freeCmd).Once()
|
||||
mockClient.On("Eval", mock.AnythingOfType("[]interface {}")).Return(premiumCmd).Once()
|
||||
|
||||
// Create multi-tier rate limiter
|
||||
tiers := map[string]struct {
|
||||
Limit int64
|
||||
Window time.Duration
|
||||
Prefix string
|
||||
}{
|
||||
"free": {Limit: 10, Window: time.Minute, Prefix: "free"},
|
||||
"premium": {Limit: 100, Window: time.Minute, Prefix: "premium"},
|
||||
}
|
||||
|
||||
multiLimiter := NewMultiTierRateLimiter(&redis.Client{}, tiers)
|
||||
|
||||
// Replace client with mock for testing
|
||||
for tier, limiter := range multiLimiter.limiters {
|
||||
if tier == "free" {
|
||||
limiter.client = mockClient
|
||||
} else if tier == "premium" {
|
||||
limiter.client = mockClient
|
||||
}
|
||||
}
|
||||
|
||||
// Test free tier
|
||||
result, err := multiLimiter.CheckTier(context.Background(), "free", "user123")
|
||||
assert.NoError(t, err)
|
||||
assert.True(t, result.Allowed)
|
||||
assert.Equal(t, int64(10), result.TotalLimit)
|
||||
|
||||
// Test premium tier
|
||||
result, err = multiLimiter.CheckTier(context.Background(), "premium", "user456")
|
||||
assert.NoError(t, err)
|
||||
assert.True(t, result.Allowed)
|
||||
assert.Equal(t, int64(100), result.TotalLimit)
|
||||
}
|
||||
|
||||
func TestRedisRateLimiter_Reset(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
mockClient := new(MockRedisClient)
|
||||
cmd := redis.NewIntCmd(context.Background())
|
||||
cmd.SetVal(1)
|
||||
mockClient.On("Del", "test:user123").Return(cmd)
|
||||
|
||||
limiter := NewRedisRateLimiter(&redis.Client{}, "test", 10, time.Minute)
|
||||
limiter.client = mockClient
|
||||
|
||||
err := limiter.Reset(context.Background(), "user123")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestRedisRateLimiter_GetStats(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
mockClient := new(MockRedisClient)
|
||||
cmd := redis.NewIntCmd(context.Background())
|
||||
cmd.SetVal(3)
|
||||
mockClient.On("ZCount", "test:user123", mock.AnythingOfType("string"), mock.AnythingOfType("string")).Return(cmd)
|
||||
|
||||
limiter := NewRedisRateLimiter(&redis.Client{}, "test", 10, time.Minute)
|
||||
limiter.client = mockClient
|
||||
|
||||
result, err := limiter.GetStats(context.Background(), "user123")
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, int64(10), result.TotalLimit)
|
||||
assert.Equal(t, int64(7), result.Remaining) // 10 - 3 = 7 remaining
|
||||
}
|
||||
Reference in New Issue
Block a user