101 lines
2.6 KiB
Go
101 lines
2.6 KiB
Go
package admission
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestFixedWindowEnforcesGlobalAndPerKeyLimits(t *testing.T) {
|
|
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
limiter, err := NewFixedWindow(FixedWindowConfig{
|
|
Window: time.Minute,
|
|
Global: 3,
|
|
PerKey: 2,
|
|
Now: func() time.Time { return now },
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewFixedWindow(): %v", err)
|
|
}
|
|
|
|
if err := limiter.Admit(context.Background(), "client-a"); err != nil {
|
|
t.Fatalf("Admit(a1): %v", err)
|
|
}
|
|
if err := limiter.Admit(context.Background(), "client-a"); err != nil {
|
|
t.Fatalf("Admit(a2): %v", err)
|
|
}
|
|
if err := limiter.Admit(context.Background(), "client-a"); !errors.Is(err, ErrPerKeyLimit) {
|
|
t.Fatalf("Admit(a3) error = %v, want ErrPerKeyLimit", err)
|
|
}
|
|
if err := limiter.Admit(context.Background(), "client-b"); err != nil {
|
|
t.Fatalf("Admit(b1): %v", err)
|
|
}
|
|
if err := limiter.Admit(context.Background(), "client-c"); !errors.Is(err, ErrGlobalLimit) {
|
|
t.Fatalf("Admit(c1) error = %v, want ErrGlobalLimit", err)
|
|
}
|
|
}
|
|
|
|
func TestFixedWindowResetsAndIsConcurrencySafe(t *testing.T) {
|
|
var current atomic.Int64
|
|
base := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
limiter, err := NewFixedWindow(FixedWindowConfig{
|
|
Window: time.Minute,
|
|
Global: 10,
|
|
PerKey: 10,
|
|
Now: func() time.Time {
|
|
return base.Add(time.Duration(current.Load()))
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewFixedWindow(): %v", err)
|
|
}
|
|
|
|
var accepted atomic.Int64
|
|
var wg sync.WaitGroup
|
|
for range 100 {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
if limiter.Admit(context.Background(), "client-a") == nil {
|
|
accepted.Add(1)
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
if got := accepted.Load(); got != 10 {
|
|
t.Fatalf("accepted = %d, want 10", got)
|
|
}
|
|
|
|
current.Store(int64(time.Minute))
|
|
if err := limiter.Admit(context.Background(), "client-a"); err != nil {
|
|
t.Fatalf("Admit(after reset): %v", err)
|
|
}
|
|
}
|
|
|
|
func TestFixedWindowRejectsNilContextAndEmptyIdentity(t *testing.T) {
|
|
t.Parallel()
|
|
limiter, err := NewFixedWindow(FixedWindowConfig{
|
|
Window: time.Minute,
|
|
Global: 1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewFixedWindow(): %v", err)
|
|
}
|
|
|
|
if err := limiter.Admit(nil, "client-a"); !errors.Is(err, ErrInvalidIdentity) {
|
|
t.Fatalf("Admit(nil context) error = %v, want ErrInvalidIdentity", err)
|
|
}
|
|
if err := limiter.Admit(context.Background(), ""); !errors.Is(err, ErrInvalidIdentity) {
|
|
t.Fatalf("Admit(empty identity) error = %v, want ErrInvalidIdentity", err)
|
|
}
|
|
}
|
|
|
|
func TestAdmitterImplementations(t *testing.T) {
|
|
t.Parallel()
|
|
var _ Admitter = (*FixedWindow)(nil)
|
|
var _ Admitter = AllowAll{}
|
|
}
|