onasty/internal/transport/http/ratelimit/ratelimit_test.go(view raw)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
package ratelimit
import (
"net/http"
"net/http/httptest"
"testing"
"testing/synctest"
"time"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
)
func TestRateLimiter_getVisitor(t *testing.T) {
limiter := newLimiter(10, 20, time.Second)
ip := visitorIP("127.0.0.1")
visitor := limiter.getVisitor(ip)
assert.NotNil(t, visitor)
visitorAgain := limiter.getVisitor(ip)
assert.Equal(t, visitor, visitorAgain)
assert.Len(t, limiter.visitors, 1)
}
func TestRateLimiter_cleanupVisitors(t *testing.T) {
synctest.Test(t, func(t *testing.T) {
limiter := newLimiter(10, 20, time.Minute)
limiter.getVisitor("192.168.9.1")
assert.Len(t, limiter.visitors, 1)
time.Sleep(61 * time.Second)
limiter.cleanupVisitors()
assert.Empty(t, limiter.visitors)
})
}
func TestMiddleware(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := map[string]struct {
config Config
requests int
expectedCode int
}{
"allows requests with in limit": {
config: Config{
RPS: 2,
Burst: 2,
TTL: time.Minute,
},
requests: 1,
expectedCode: http.StatusOK,
},
"blocks requests over limit": {
config: Config{
RPS: 1,
Burst: 1,
TTL: time.Minute,
},
requests: 2,
expectedCode: http.StatusTooManyRequests,
},
"allows burst requests": {
config: Config{
RPS: 1,
Burst: 3,
TTL: time.Minute,
},
requests: 3,
expectedCode: http.StatusOK,
},
}
for name, tt := range tests {
t.Run(name, func(t *testing.T) {
handler := MiddlewareWithConfig(tt.config)
var lastCode int
for range tt.requests {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
handler(c)
lastCode = w.Code
}
assert.Equal(t, tt.expectedCode, lastCode)
})
}
}
|