56 lines
953 B
Go
56 lines
953 B
Go
|
package httpapi
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
"sync/atomic"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestSlidingWindowLimiterAllowConcurrent(t *testing.T) {
|
||
|
|
t.Parallel()
|
||
|
|
l := newSlidingWindowLimiter(15, time.Minute)
|
||
|
|
var allowed atomic.Int64
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
const n = 50
|
||
|
|
wg.Add(n)
|
||
|
|
for i := 0; i < n; i++ {
|
||
|
|
go func() {
|
||
|
|
defer wg.Done()
|
||
|
|
if l.allow("company-a") {
|
||
|
|
allowed.Add(1)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
if got := allowed.Load(); got != 15 {
|
||
|
|
t.Fatalf("allowed=%d want 15", got)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestSlidingWindowLimiterSeparateKeys(t *testing.T) {
|
||
|
|
t.Parallel()
|
||
|
|
l := newSlidingWindowLimiter(5, time.Minute)
|
||
|
|
var a, b atomic.Int64
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
wg.Add(20)
|
||
|
|
for i := 0; i < 10; i++ {
|
||
|
|
go func() {
|
||
|
|
defer wg.Done()
|
||
|
|
if l.allow("a") {
|
||
|
|
a.Add(1)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
go func() {
|
||
|
|
defer wg.Done()
|
||
|
|
if l.allow("b") {
|
||
|
|
b.Add(1)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
if a.Load() != 5 || b.Load() != 5 {
|
||
|
|
t.Fatalf("a=%d b=%d want 5 each", a.Load(), b.Load())
|
||
|
|
}
|
||
|
|
}
|