traefik/pkg/middlewares/ratelimiter/rate_limiter_test.go
mpl 4ec90c5c0d Add rate limiter, rename maxConn into inFlightReq
Co-authored-by: Julien Salleyron <julien.salleyron@gmail.com>
Co-authored-by: Jean-Baptiste Doumenjou <jb.doumenjou@gmail.com>
2019-08-26 03:20:06 -07:00

160 lines
3.6 KiB
Go

package ratelimiter
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/containous/traefik/v2/pkg/config/dynamic"
"github.com/containous/traefik/v2/pkg/testhelpers"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/vulcand/oxy/utils"
)
func TestNewRateLimiter(t *testing.T) {
testCases := []struct {
desc string
config dynamic.RateLimit
expectedMaxDelay time.Duration
expectedSourceIP string
}{
{
desc: "maxDelay computation",
config: dynamic.RateLimit{
Average: 200,
Burst: 10,
},
expectedMaxDelay: 2500 * time.Microsecond,
},
{
desc: "default SourceMatcher is remote address ip strategy",
config: dynamic.RateLimit{
Average: 200,
Burst: 10,
},
expectedSourceIP: "127.0.0.1",
},
}
for _, test := range testCases {
test := test
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, test.config, "rate-limiter")
require.NoError(t, err)
rtl, _ := h.(*rateLimiter)
if test.expectedMaxDelay != 0 {
assert.Equal(t, test.expectedMaxDelay, rtl.maxDelay)
}
if test.expectedSourceIP != "" {
extractor, ok := rtl.sourceMatcher.(utils.ExtractorFunc)
require.True(t, ok, "Not an ExtractorFunc")
req := http.Request{
RemoteAddr: fmt.Sprintf("%s:1234", test.expectedSourceIP),
}
ip, _, err := extractor(&req)
assert.NoError(t, err)
assert.Equal(t, test.expectedSourceIP, ip)
}
})
}
}
func TestRateLimit(t *testing.T) {
testCases := []struct {
desc string
config dynamic.RateLimit
reqCount int
}{
{
desc: "Average is respected",
config: dynamic.RateLimit{
Average: 100,
Burst: 1,
},
reqCount: 200,
},
{
desc: "Burst is taken into account",
config: dynamic.RateLimit{
Average: 100,
Burst: 200,
},
reqCount: 300,
},
{
desc: "Zero average ==> no rate limiting",
config: dynamic.RateLimit{
Average: 0,
Burst: 1,
},
reqCount: 100,
},
}
for _, test := range testCases {
test := test
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
reqCount := 0
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
reqCount++
})
h, err := New(context.Background(), next, test.config, "rate-limiter")
require.NoError(t, err)
start := time.Now()
for {
if reqCount >= test.reqCount {
break
}
req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil)
req.RemoteAddr = "127.0.0.1:1234"
w := httptest.NewRecorder()
h.ServeHTTP(w, req)
// TODO(mpl): predict and count the 200 VS the 429?
}
stop := time.Now()
elapsed := stop.Sub(start)
if test.config.Average == 0 {
if elapsed > time.Millisecond {
t.Fatalf("rate should not have been limited, but: %d requests in %v", reqCount, elapsed)
}
return
}
// Assume allowed burst is initially consumed in an infinitesimal period of time
var expectedDuration time.Duration
if test.config.Average != 0 {
expectedDuration = time.Duration((int64(test.reqCount)-test.config.Burst+1)/test.config.Average) * time.Second
}
// Allow for a 2% leeway
minDuration := expectedDuration * 98 / 100
maxDuration := expectedDuration * 102 / 100
if elapsed < minDuration {
t.Fatalf("rate was faster than expected: %d requests in %v", reqCount, elapsed)
}
if elapsed > maxDuration {
t.Fatalf("rate was slower than expected: %d requests in %v", reqCount, elapsed)
}
})
}
}