telecom-platform/apps/api-server/internal/middleware/redis_ratelimit_test.go

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 {
	argsList := make([]interface{}, 0, len(keys)+len(args)+1)
	argsList = append(argsList, script)
	for _, key := range keys {
		argsList = append(argsList, key)
	}
	argsList = append(argsList, args...)
	return m.Called(argsList...).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)
			// Match the actual call signature: script, key, now, window, limit
			mockClient.On("Eval", mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("int64"), mock.AnythingOfType("int64"), mock.AnythingOfType("int64")).Return(cmd)

			// Create rate limiter with mock client
			limiter := NewRedisRateLimiter(mockClient, "test", tt.limit, tt.window)

			// 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("string"), mock.AnythingOfType("string"), mock.AnythingOfType("int64"), mock.AnythingOfType("int64"), mock.AnythingOfType("int64")).Return(cmd)

			// Create rate limiter with mock client
			limiter := NewRedisRateLimiter(mockClient, "test", 5, time.Minute)

			// 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   any
		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("string"), mock.AnythingOfType("string"), mock.AnythingOfType("int64"), mock.AnythingOfType("int64"), mock.AnythingOfType("int64")).Return(freeCmd).Once()
	mockClient.On("Eval", mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("int64"), mock.AnythingOfType("int64"), mock.AnythingOfType("int64")).Return(premiumCmd).Once()

	// Create multi-tier rate limiter with mock client
	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(mockClient, tiers)

	// 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(mockClient, "test", 10, time.Minute)

	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(mockClient, "test", 10, time.Minute)

	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
}