Lithium/middlewares/ratelimiter_test.go

package middlewares

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

	"github.com/stretchr/testify/assert"
)

func ExecuteRequest(middlewareHandler http.Handler) int {
	request, _ := http.NewRequest("GET", "/", nil)
	responseRecorder := httptest.NewRecorder()

	middlewareHandler.ServeHTTP(responseRecorder, request)

	return responseRecorder.Code
}

func TestRateLimiterMiddleware(t *testing.T) {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	t.Run("AuthorizationMiddleware returns 200 response when rate limit is not hit", func(t *testing.T) {
		middleware, err := CreateRateLimiterMiddleware(1, 0)
		assert.Nil(t, err)

		middlewareHandler := middleware.Middleware(handler)

		assert.Equal(t, 200, ExecuteRequest(middlewareHandler))
	})

	t.Run("AuthorizationMiddleware returns 429 response when rate limit is hit", func(t *testing.T) {
		middleware, err := CreateRateLimiterMiddleware(1, 0)
		assert.Nil(t, err)

		middlewareHandler := middleware.Middleware(handler)

		assert.Equal(t, 200, ExecuteRequest(middlewareHandler))
		assert.Equal(t, 429, ExecuteRequest(middlewareHandler))
	})

	t.Run("AuthorizationMiddleware returns 200 response when rate limit with burst is not hit", func(t *testing.T) {
		middleware, err := CreateRateLimiterMiddleware(1, 1)
		assert.Nil(t, err)

		middlewareHandler := middleware.Middleware(handler)

		assert.Equal(t, 200, ExecuteRequest(middlewareHandler))
		assert.Equal(t, 200, ExecuteRequest(middlewareHandler))
	})

	t.Run("AuthorizationMiddleware returns 429 response when rate limit with burst is hit", func(t *testing.T) {
		middleware, err := CreateRateLimiterMiddleware(1, 1)
		assert.Nil(t, err)

		middlewareHandler := middleware.Middleware(handler)

		assert.Equal(t, 200, ExecuteRequest(middlewareHandler))
		assert.Equal(t, 200, ExecuteRequest(middlewareHandler))
		assert.Equal(t, 429, ExecuteRequest(middlewareHandler))
	})
}