Lithium/middlewares/ratelimiter_test.go

package middlewares

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

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

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)
		assert.Nil(t, err)

		middlewareHandler := middleware.Middleware(handler)

		request, _ := http.NewRequest("GET", "/", nil)
		responseRecorder := httptest.NewRecorder()

		middlewareHandler.ServeHTTP(responseRecorder, request)

		assert.Equal(t, 200, responseRecorder.Code)
	})

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

		middlewareHandler := middleware.Middleware(handler)

		request, _ := http.NewRequest("GET", "/", nil)
		responseRecorder := httptest.NewRecorder()

		middlewareHandler.ServeHTTP(responseRecorder, request)
		assert.Equal(t, 200, responseRecorder.Code)

		request, _ = http.NewRequest("GET", "/", nil)
		responseRecorder = httptest.NewRecorder()

		middlewareHandler.ServeHTTP(responseRecorder, request)
		assert.Equal(t, 429, responseRecorder.Code)
	})
}