77 lines
2.5 KiB
Go
77 lines
2.5 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestCombinedBucketsRejectWithoutPartialConsumption(t *testing.T) {
|
|
limiter := New()
|
|
now := time.Unix(1_700_000_000, 0).UTC()
|
|
user := Bucket{Key: "user:1", Policy: Policy{Capacity: 2, Window: time.Minute}}
|
|
key := Bucket{Key: "key:1", Policy: Policy{Capacity: 1, Window: time.Minute}}
|
|
if decision := limiter.Take(now, user, key); !decision.Allowed {
|
|
t.Fatal("first combined request was rejected")
|
|
}
|
|
decision := limiter.Take(now.Add(time.Second), user, key)
|
|
if decision.Allowed || !decision.RetryAt.Equal(now.Add(time.Minute)) {
|
|
t.Fatalf("second combined request = %#v", decision)
|
|
}
|
|
if decision := limiter.Take(now.Add(2*time.Second), user); !decision.Allowed {
|
|
t.Fatal("rejected combined request consumed the user bucket")
|
|
}
|
|
if decision := limiter.Take(now.Add(time.Minute), user, key); !decision.Allowed {
|
|
t.Fatal("new fixed window was rejected")
|
|
}
|
|
}
|
|
|
|
func TestReservationCancelRefundsOnlyCurrentWindow(t *testing.T) {
|
|
limiter := New()
|
|
policy := Policy{Capacity: 1, Window: time.Second}
|
|
now := time.Unix(1_700_000_000, 0).UTC()
|
|
decision := limiter.Take(now, Bucket{Key: "provider:1", Policy: policy})
|
|
if !decision.Allowed || !decision.Reservation().Cancel() {
|
|
t.Fatal("provider reservation was not canceled")
|
|
}
|
|
if decision.Reservation().Cancel() {
|
|
t.Fatal("provider reservation canceled twice")
|
|
}
|
|
if next := limiter.Take(now, Bucket{Key: "provider:1", Policy: policy}); !next.Allowed {
|
|
t.Fatal("canceled reservation was not refunded")
|
|
}
|
|
}
|
|
|
|
func TestConcurrentTakeNeverExceedsCapacity(t *testing.T) {
|
|
limiter := New()
|
|
now := time.Unix(1_700_000_000, 0).UTC()
|
|
bucket := Bucket{Key: "provider:1", Policy: Policy{Capacity: 10, Window: time.Minute}}
|
|
var allowed atomic.Int32
|
|
var wait sync.WaitGroup
|
|
for range 100 {
|
|
wait.Add(1)
|
|
go func() {
|
|
defer wait.Done()
|
|
if limiter.Take(now, bucket).Allowed {
|
|
allowed.Add(1)
|
|
}
|
|
}()
|
|
}
|
|
wait.Wait()
|
|
if got := allowed.Load(); got != 10 {
|
|
t.Fatalf("allowed = %d, want 10", got)
|
|
}
|
|
}
|
|
|
|
func TestTakeRejectsDuplicateBucketKeys(t *testing.T) {
|
|
limiter := New()
|
|
bucket := Bucket{Key: "same", Policy: Policy{Capacity: 1, Window: time.Minute}}
|
|
if decision := limiter.Take(time.Now().UTC(), bucket, bucket); decision.Allowed || !decision.RetryAt.IsZero() {
|
|
t.Fatalf("duplicate bucket decision = %#v", decision)
|
|
}
|
|
if decision := limiter.Take(time.Now().UTC(), bucket); !decision.Allowed {
|
|
t.Fatal("invalid duplicate request consumed capacity")
|
|
}
|
|
}
|