mirror of
https://github.com/empayre/fleet.git
synced 2024-11-06 17:05:18 +00:00
35 lines
882 B
Go
35 lines
882 B
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/throttled/throttled/v2"
|
|
"github.com/throttled/throttled/v2/store/memstore"
|
|
)
|
|
|
|
// Intent is to test the middleware functionality. We rely on the tests within
|
|
// Throttled to verify that the rate limiting algorithm works properly.
|
|
|
|
func TestLimit(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
store, _ := memstore.New(0)
|
|
limiter := NewMiddleware(store)
|
|
endpoint := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
|
|
wrapped := limiter.Limit(
|
|
throttled.RateQuota{MaxRate: throttled.PerHour(1), MaxBurst: 0},
|
|
)(endpoint)
|
|
|
|
_, err := wrapped(context.Background(), struct{}{})
|
|
assert.NoError(t, err)
|
|
|
|
// Hits rate limit
|
|
_, err = wrapped(context.Background(), struct{}{})
|
|
assert.Error(t, err)
|
|
var rle Error
|
|
assert.True(t, errors.As(err, &rle))
|
|
}
|