f737002baa
Integrates the custom rate limiter in the previous PR into the rest of the application. Also takes the opportunity to clean up how we instantiate Redis clients so that we can share them among multiple different services. There are some config changes in this PR. Specifically, you must specify a `base_rate` and `base_interval` in the rate limit config.
54 lines
1.2 KiB
Go
54 lines
1.2 KiB
Go
package proxyd
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/alicebob/miniredis"
|
|
"github.com/go-redis/redis/v8"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestFrontendRateLimiter(t *testing.T) {
|
|
redisServer, err := miniredis.Run()
|
|
require.NoError(t, err)
|
|
defer redisServer.Close()
|
|
|
|
redisClient := redis.NewClient(&redis.Options{
|
|
Addr: fmt.Sprintf("127.0.0.1:%s", redisServer.Port()),
|
|
})
|
|
|
|
max := 2
|
|
lims := []struct {
|
|
name string
|
|
frl FrontendRateLimiter
|
|
}{
|
|
{"memory", NewMemoryFrontendRateLimit(2*time.Second, max)},
|
|
{"redis", NewRedisFrontendRateLimiter(redisClient, 2*time.Second, max, "")},
|
|
}
|
|
|
|
for _, cfg := range lims {
|
|
frl := cfg.frl
|
|
ctx := context.Background()
|
|
t.Run(cfg.name, func(t *testing.T) {
|
|
for i := 0; i < 4; i++ {
|
|
ok, err := frl.Take(ctx, "foo")
|
|
require.NoError(t, err)
|
|
require.Equal(t, i < max, ok)
|
|
ok, err = frl.Take(ctx, "bar")
|
|
require.NoError(t, err)
|
|
require.Equal(t, i < max, ok)
|
|
}
|
|
time.Sleep(2 * time.Second)
|
|
for i := 0; i < 4; i++ {
|
|
ok, _ := frl.Take(ctx, "foo")
|
|
require.Equal(t, i < max, ok)
|
|
ok, _ = frl.Take(ctx, "bar")
|
|
require.Equal(t, i < max, ok)
|
|
}
|
|
})
|
|
}
|
|
}
|