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 }