133 lines
3.8 KiB
Go
133 lines
3.8 KiB
Go
package ratelimit_test
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/zap"
|
|
|
|
"github.com/veylant/ia-gateway/internal/ratelimit"
|
|
)
|
|
|
|
func newTestLimiter(tenantRPM, tenantBurst, userRPM, userBurst int) *ratelimit.Limiter {
|
|
cfg := ratelimit.RateLimitConfig{
|
|
RequestsPerMin: tenantRPM,
|
|
BurstSize: tenantBurst,
|
|
UserRPM: userRPM,
|
|
UserBurst: userBurst,
|
|
IsEnabled: true,
|
|
}
|
|
return ratelimit.New(cfg, zap.NewNop())
|
|
}
|
|
|
|
func TestAllow_WithinLimits(t *testing.T) {
|
|
l := newTestLimiter(600, 5, 300, 5)
|
|
// Should allow up to burst size immediately.
|
|
for i := 0; i < 5; i++ {
|
|
assert.True(t, l.Allow("tenant1", "user1"), "request %d should be allowed", i+1)
|
|
}
|
|
}
|
|
|
|
func TestAllow_TenantLimitExceeded(t *testing.T) {
|
|
// Very low tenant burst (1) — second request should fail.
|
|
l := newTestLimiter(600, 1, 600, 100)
|
|
assert.True(t, l.Allow("t1", "u1"))
|
|
assert.False(t, l.Allow("t1", "u1"), "second request should be blocked by tenant bucket")
|
|
}
|
|
|
|
func TestAllow_UserLimitExceeded(t *testing.T) {
|
|
// High tenant burst, low user burst (1) — second request from same user should fail.
|
|
l := newTestLimiter(600, 100, 600, 1)
|
|
assert.True(t, l.Allow("t1", "u1"))
|
|
assert.False(t, l.Allow("t1", "u1"), "second request from same user should be blocked")
|
|
}
|
|
|
|
func TestAllow_DifferentUsers_IndependentBuckets(t *testing.T) {
|
|
l := newTestLimiter(600, 100, 600, 1)
|
|
// u1 exhausts its bucket, u2 should still be allowed.
|
|
l.Allow("t1", "u1") // consumes u1's burst
|
|
assert.False(t, l.Allow("t1", "u1"), "u1 should be blocked")
|
|
assert.True(t, l.Allow("t1", "u2"), "u2 should be allowed (independent bucket)")
|
|
}
|
|
|
|
func TestAllow_Disabled_AlwaysAllowed(t *testing.T) {
|
|
cfg := ratelimit.RateLimitConfig{
|
|
TenantID: "t1",
|
|
RequestsPerMin: 1, // extremely low
|
|
BurstSize: 1,
|
|
UserRPM: 1,
|
|
UserBurst: 1,
|
|
IsEnabled: false, // disabled → bypass
|
|
}
|
|
l := ratelimit.New(cfg, zap.NewNop())
|
|
for i := 0; i < 100; i++ {
|
|
require.True(t, l.Allow("t1", "u1"), "disabled limiter must always allow")
|
|
}
|
|
}
|
|
|
|
func TestSetConfig_OverridesDefault(t *testing.T) {
|
|
// Default: burst=1 (very restrictive)
|
|
l := newTestLimiter(600, 1, 600, 1)
|
|
|
|
// Override tenant t1 with a generous limit.
|
|
l.SetConfig(ratelimit.RateLimitConfig{
|
|
TenantID: "t1",
|
|
RequestsPerMin: 6000,
|
|
BurstSize: 100,
|
|
UserRPM: 6000,
|
|
UserBurst: 100,
|
|
IsEnabled: true,
|
|
})
|
|
|
|
// Should allow many requests now.
|
|
for i := 0; i < 50; i++ {
|
|
assert.True(t, l.Allow("t1", "u1"), "overridden config should allow request %d", i+1)
|
|
}
|
|
}
|
|
|
|
func TestDeleteConfig_RevertsToDefault(t *testing.T) {
|
|
l := newTestLimiter(600, 100, 600, 100)
|
|
l.SetConfig(ratelimit.RateLimitConfig{
|
|
TenantID: "t1",
|
|
RequestsPerMin: 6000,
|
|
BurstSize: 1000,
|
|
UserRPM: 6000,
|
|
UserBurst: 1000,
|
|
IsEnabled: true,
|
|
})
|
|
l.DeleteConfig("t1")
|
|
|
|
// Default burst is 100, so many requests should still be allowed immediately.
|
|
assert.True(t, l.Allow("t1", "u1"))
|
|
}
|
|
|
|
func TestListConfigs(t *testing.T) {
|
|
l := newTestLimiter(600, 100, 600, 100)
|
|
assert.Empty(t, l.ListConfigs())
|
|
|
|
l.SetConfig(ratelimit.RateLimitConfig{TenantID: "t1", RequestsPerMin: 100, BurstSize: 10, UserRPM: 50, UserBurst: 5, IsEnabled: true})
|
|
l.SetConfig(ratelimit.RateLimitConfig{TenantID: "t2", RequestsPerMin: 200, BurstSize: 20, UserRPM: 100, UserBurst: 10, IsEnabled: true})
|
|
|
|
cfgs := l.ListConfigs()
|
|
assert.Len(t, cfgs, 2)
|
|
}
|
|
|
|
func TestConcurrentAllow_RaceFree(t *testing.T) {
|
|
l := newTestLimiter(6000, 1000, 6000, 1000)
|
|
done := make(chan struct{})
|
|
for i := 0; i < 10; i++ {
|
|
go func(id int) {
|
|
for j := 0; j < 100; j++ {
|
|
l.Allow("tenant", "user")
|
|
time.Sleep(time.Microsecond)
|
|
}
|
|
done <- struct{}{}
|
|
}(i)
|
|
}
|
|
for i := 0; i < 10; i++ {
|
|
<-done
|
|
}
|
|
}
|