82 lines
1.7 KiB
Go
82 lines
1.7 KiB
Go
|
package processing
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"sync"
|
||
|
|
"sync/atomic"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/google/uuid"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestWaitRateSerializesConcurrentCallers(t *testing.T) {
|
||
|
|
t.Parallel()
|
||
|
|
c := &OpenAIClient{MinInterval: 30 * time.Millisecond}
|
||
|
|
const n = 8
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
wg.Add(n)
|
||
|
|
start := time.Now()
|
||
|
|
for i := 0; i < n; i++ {
|
||
|
|
go func() {
|
||
|
|
defer wg.Done()
|
||
|
|
if err := c.waitRate(context.Background()); err != nil {
|
||
|
|
t.Errorf("waitRate: %v", err)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
elapsed := time.Since(start)
|
||
|
|
// With reservation under the lock, n callers need ~ (n-1)*MinInterval.
|
||
|
|
minExpected := time.Duration(n-2) * c.MinInterval
|
||
|
|
if elapsed < minExpected {
|
||
|
|
t.Fatalf("elapsed %v too short for %d serialized waits (want >= %v)", elapsed, n, minExpected)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestStartLimiterAllowConcurrent(t *testing.T) {
|
||
|
|
t.Parallel()
|
||
|
|
l := NewStartLimiter(10, time.Minute)
|
||
|
|
company := uuid.New()
|
||
|
|
var allowed atomic.Int64
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
const n = 40
|
||
|
|
wg.Add(n)
|
||
|
|
for i := 0; i < n; i++ {
|
||
|
|
go func() {
|
||
|
|
defer wg.Done()
|
||
|
|
if l.Allow(company) {
|
||
|
|
allowed.Add(1)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
if got := allowed.Load(); got != 10 {
|
||
|
|
t.Fatalf("allowed=%d want 10", got)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestStartLimiterSeparateCompanies(t *testing.T) {
|
||
|
|
t.Parallel()
|
||
|
|
l := NewStartLimiter(3, time.Minute)
|
||
|
|
a, b := uuid.New(), uuid.New()
|
||
|
|
for i := 0; i < 3; i++ {
|
||
|
|
if !l.Allow(a) {
|
||
|
|
t.Fatalf("company A start %d should allow", i)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if l.Allow(a) {
|
||
|
|
t.Fatal("company A should be rate limited")
|
||
|
|
}
|
||
|
|
if !l.Allow(b) {
|
||
|
|
t.Fatal("company B should not share A budget")
|
||
|
|
}
|
||
|
|
if NewStartLimiter(1, time.Minute) == nil {
|
||
|
|
t.Fatal("NewStartLimiter must return non-nil")
|
||
|
|
}
|
||
|
|
var nilLimiter *StartLimiter
|
||
|
|
if !nilLimiter.Allow(a) {
|
||
|
|
t.Fatal("nil StartLimiter must allow (fail-open)")
|
||
|
|
}
|
||
|
|
}
|