github.com/adharshmk96/stk@v1.2.3/pkg/middleware/rate_limiter_test.go (about)

     1  package middleware_test
     2  
     3  import (
     4  	"net/http"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/adharshmk96/stk/gsk"
     9  	"github.com/adharshmk96/stk/pkg/middleware"
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  func dummyHandler(c *gsk.Context) {
    14  	c.Status(http.StatusOK).JSONResponse("OK")
    15  }
    16  
    17  func TestRateLimiter(t *testing.T) {
    18  	// Create a new server instance
    19  	config := &gsk.ServerConfig{
    20  		Port: "8888",
    21  	}
    22  	s := gsk.New(config)
    23  
    24  	// rate limiter middleware
    25  	requestsPerInterval := 5
    26  	interval := 1 * time.Second
    27  
    28  	rlConfig := middleware.RateLimiterConfig{
    29  		RequestsPerInterval: requestsPerInterval,
    30  		Interval:            interval,
    31  	}
    32  
    33  	rateLimiter := middleware.NewRateLimiter(rlConfig)
    34  
    35  	s.Use(rateLimiter.Middleware)
    36  
    37  	s.Get("/test", dummyHandler)
    38  
    39  	for i := 0; i < requestsPerInterval; i++ {
    40  		rr, _ := s.Test("GET", "/test", nil)
    41  
    42  		if rr.Code != http.StatusOK {
    43  			t.Errorf("Expected 200 OK, got: %d for request %d", rr.Code, i+1)
    44  		}
    45  	}
    46  
    47  	rr, _ := s.Test("GET", "/test", nil)
    48  
    49  	assert.Equal(t, http.StatusTooManyRequests, rr.Code)
    50  }