veza/veza-backend-api/internal/middleware/ratelimit_test.go

package middleware

import (
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestSimpleRateLimiter_WithinLimit(t *testing.T) {
	gin.SetMode(gin.TestMode)
	limiter := NewSimpleRateLimiter(5, 1*time.Minute)

	router := gin.New()
	router.Use(limiter.Middleware())
	router.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"ok": true})
	})

	// Faire 5 requêtes (dans la limite)
	for i := 0; i < 5; i++ {
		w := httptest.NewRecorder()
		req := httptest.NewRequest("GET", "/test", nil)
		req.RemoteAddr = "127.0.0.1:12345"
		router.ServeHTTP(w, req)

		assert.Equal(t, http.StatusOK, w.Code)
		assert.Equal(t, "5", w.Header().Get("X-RateLimit-Limit"))
		remaining := w.Header().Get("X-RateLimit-Remaining")
		assert.NotEmpty(t, remaining)
		assert.Contains(t, []string{"4", "3", "2", "1", "0"}, remaining)
	}
}

func TestSimpleRateLimiter_ExceedsLimit(t *testing.T) {
	gin.SetMode(gin.TestMode)
	limiter := NewSimpleRateLimiter(5, 1*time.Minute)

	router := gin.New()
	router.Use(limiter.Middleware())
	router.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"ok": true})
	})

	// Faire 5 requêtes (dans la limite)
	for i := 0; i < 5; i++ {
		w := httptest.NewRecorder()
		req := httptest.NewRequest("GET", "/test", nil)
		req.RemoteAddr = "127.0.0.1:12345"
		router.ServeHTTP(w, req)
		assert.Equal(t, http.StatusOK, w.Code)
	}

	// 6ème requête devrait être bloquée
	w := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/test", nil)
	req.RemoteAddr = "127.0.0.1:12345"
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusTooManyRequests, w.Code)
	assert.Equal(t, "5", w.Header().Get("X-RateLimit-Limit"))
	assert.Equal(t, "0", w.Header().Get("X-RateLimit-Remaining"))
	assert.NotEmpty(t, w.Header().Get("X-RateLimit-Reset"))
}

func TestSimpleRateLimiter_DifferentIPs(t *testing.T) {
	gin.SetMode(gin.TestMode)
	limiter := NewSimpleRateLimiter(5, 1*time.Minute)

	router := gin.New()
	router.Use(limiter.Middleware())
	router.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"ok": true})
	})

	// IP 1: 5 requêtes (dans la limite)
	for i := 0; i < 5; i++ {
		w := httptest.NewRecorder()
		req := httptest.NewRequest("GET", "/test", nil)
		req.RemoteAddr = "127.0.0.1:12345"
		router.ServeHTTP(w, req)
		assert.Equal(t, http.StatusOK, w.Code)
	}

	// IP 2: 5 requêtes (devrait aussi être dans la limite car IP différente)
	for i := 0; i < 5; i++ {
		w := httptest.NewRecorder()
		req := httptest.NewRequest("GET", "/test", nil)
		req.RemoteAddr = "192.168.1.1:12345"
		router.ServeHTTP(w, req)
		assert.Equal(t, http.StatusOK, w.Code, "IP différente devrait avoir sa propre limite")
	}
}

func TestSimpleRateLimiter_Headers(t *testing.T) {
	gin.SetMode(gin.TestMode)
	limiter := NewSimpleRateLimiter(100, 1*time.Minute)

	router := gin.New()
	router.Use(limiter.Middleware())
	router.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"ok": true})
	})

	w := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/test", nil)
	req.RemoteAddr = "127.0.0.1:12345"
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)
	assert.Equal(t, "100", w.Header().Get("X-RateLimit-Limit"))
	assert.Equal(t, "99", w.Header().Get("X-RateLimit-Remaining"))
	assert.NotEmpty(t, w.Header().Get("X-RateLimit-Reset"))
}

func TestSimpleRateLimiter_WindowExpiration(t *testing.T) {
	gin.SetMode(gin.TestMode)
	// Utiliser une fenêtre très courte pour les tests
	limiter := NewSimpleRateLimiter(2, 100*time.Millisecond)

	router := gin.New()
	router.Use(limiter.Middleware())
	router.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"ok": true})
	})

	// Faire 2 requêtes (limite atteinte)
	w1 := httptest.NewRecorder()
	req1 := httptest.NewRequest("GET", "/test", nil)
	req1.RemoteAddr = "127.0.0.1:12345"
	router.ServeHTTP(w1, req1)
	assert.Equal(t, http.StatusOK, w1.Code)

	w2 := httptest.NewRecorder()
	req2 := httptest.NewRequest("GET", "/test", nil)
	req2.RemoteAddr = "127.0.0.1:12345"
	router.ServeHTTP(w2, req2)
	assert.Equal(t, http.StatusOK, w2.Code)

	// 3ème requête devrait être bloquée
	w3 := httptest.NewRecorder()
	req3 := httptest.NewRequest("GET", "/test", nil)
	req3.RemoteAddr = "127.0.0.1:12345"
	router.ServeHTTP(w3, req3)
	assert.Equal(t, http.StatusTooManyRequests, w3.Code)

	// Attendre que la fenêtre expire
	time.Sleep(150 * time.Millisecond)

	// Après expiration, une nouvelle requête devrait passer
	w4 := httptest.NewRecorder()
	req4 := httptest.NewRequest("GET", "/test", nil)
	req4.RemoteAddr = "127.0.0.1:12345"
	router.ServeHTTP(w4, req4)
	assert.Equal(t, http.StatusOK, w4.Code, "Après expiration de la fenêtre, la requête devrait passer")
}

func TestNewSimpleRateLimiter(t *testing.T) {
	limiter := NewSimpleRateLimiter(100, 1*time.Minute)
	require.NotNil(t, limiter)
	assert.Equal(t, 100, limiter.limit)
	assert.Equal(t, 1*time.Minute, limiter.window)
	assert.NotNil(t, limiter.requests)

	// Arrêter la goroutine de nettoyage
	// (dans un vrai test, on pourrait ajouter une méthode Stop())
}

func TestSimpleRateLimiter_ErrorResponse(t *testing.T) {
	gin.SetMode(gin.TestMode)
	limiter := NewSimpleRateLimiter(1, 1*time.Minute)

	router := gin.New()
	router.Use(limiter.Middleware())
	router.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"ok": true})
	})

	// Première requête OK
	w1 := httptest.NewRecorder()
	req1 := httptest.NewRequest("GET", "/test", nil)
	req1.RemoteAddr = "127.0.0.1:12345"
	router.ServeHTTP(w1, req1)
	assert.Equal(t, http.StatusOK, w1.Code)

	// Deuxième requête bloquée
	w2 := httptest.NewRecorder()
	req2 := httptest.NewRequest("GET", "/test", nil)
	req2.RemoteAddr = "127.0.0.1:12345"
	router.ServeHTTP(w2, req2)

	assert.Equal(t, http.StatusTooManyRequests, w2.Code)
	// Vérifier que le body contient le message d'erreur
	assert.Contains(t, w2.Body.String(), "Rate limit exceeded")
	assert.Contains(t, w2.Body.String(), "retry_after")
}

func TestSimpleRateLimiter_RemainingHeader(t *testing.T) {
	gin.SetMode(gin.TestMode)
	limiter := NewSimpleRateLimiter(10, 1*time.Minute)

	router := gin.New()
	router.Use(limiter.Middleware())
	router.GET("/test", func(c *gin.Context) {
		c.JSON(200, gin.H{"ok": true})
	})

	// Faire plusieurs requêtes et vérifier que le header Remaining diminue
	for i := 0; i < 5; i++ {
		w := httptest.NewRecorder()
		req := httptest.NewRequest("GET", "/test", nil)
		req.RemoteAddr = "127.0.0.1:12345"
		router.ServeHTTP(w, req)

		expectedRemaining := 10 - (i + 1)
		assert.Equal(t, strconv.Itoa(expectedRemaining), w.Header().Get("X-RateLimit-Remaining"))
	}
}