ginuerzh
136 lines
3.1 KiB
Go
136 lines
3.1 KiB
Go
package traffic
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
limiter "github.com/go-gost/core/limiter/traffic"
|
|
)
|
|
|
|
func TestNewLimiter(t *testing.T) {
|
|
l := NewLimiter(100)
|
|
if l == nil {
|
|
t.Fatal("NewLimiter should not be nil")
|
|
}
|
|
if l.Limit() != 100 {
|
|
t.Fatalf("expected limit 100, got %d", l.Limit())
|
|
}
|
|
}
|
|
|
|
func TestLimiter_Wait(t *testing.T) {
|
|
// High burst so Wait does not block.
|
|
l := NewLimiter(1000000)
|
|
n := l.Wait(context.Background(), 3)
|
|
if n != 3 {
|
|
t.Fatalf("expected 3, got %d", n)
|
|
}
|
|
}
|
|
|
|
func TestLimiter_Wait_BurstClamped(t *testing.T) {
|
|
// Burst is 5, request 100 -> clamped to burst.
|
|
l := NewLimiter(5)
|
|
n := l.Wait(context.Background(), 100)
|
|
if n != 5 {
|
|
t.Fatalf("expected burst 5, got %d", n)
|
|
}
|
|
}
|
|
|
|
func TestLimiter_Wait_ContextCancelled(t *testing.T) {
|
|
l := NewLimiter(1000)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
n := l.Wait(ctx, 10)
|
|
if n != 0 {
|
|
t.Fatalf("cancelled context should return 0, got %d", n)
|
|
}
|
|
}
|
|
|
|
func TestLimiter_Set(t *testing.T) {
|
|
l := NewLimiter(5)
|
|
l.Set(10)
|
|
if l.Limit() != 10 {
|
|
t.Fatalf("expected limit 10 after Set, got %d", l.Limit())
|
|
}
|
|
// Verify burst also updated: Wait should allow up to 10.
|
|
n := l.Wait(context.Background(), 10)
|
|
if n != 10 {
|
|
t.Fatalf("expected burst 10 after Set, got %d", n)
|
|
}
|
|
}
|
|
|
|
func TestLimiter_String(t *testing.T) {
|
|
l := NewLimiter(42)
|
|
// llimiter implements fmt.Stringer (not part of core interface).
|
|
if s := l.(*llimiter).String(); s != "42" {
|
|
t.Fatalf("expected '42', got '%s'", s)
|
|
}
|
|
}
|
|
|
|
func TestLimiterGroup_SortedByLimit(t *testing.T) {
|
|
g := newLimiterGroup(
|
|
NewLimiter(200),
|
|
NewLimiter(50),
|
|
NewLimiter(100),
|
|
)
|
|
if g.Limit() != 50 {
|
|
t.Fatalf("expected smallest limit 50, got %d", g.Limit())
|
|
}
|
|
}
|
|
|
|
func TestLimiterGroup_Wait(t *testing.T) {
|
|
// Group: [3, 100000] — smallest limit is 3.
|
|
small := NewLimiter(3)
|
|
large := NewLimiter(100000)
|
|
g := newLimiterGroup(large, small)
|
|
|
|
// Wait for 3: small returns 3, large returns 3, min = 3.
|
|
n := g.Wait(context.Background(), 3)
|
|
if n != 3 {
|
|
t.Fatalf("expected 3, got %d", n)
|
|
}
|
|
|
|
// Wait for 5: small clamps to burst=3, returns 3, large returns 5, min = 3.
|
|
n = g.Wait(context.Background(), 5)
|
|
if n != 3 {
|
|
t.Fatalf("expected min=3, got %d", n)
|
|
}
|
|
}
|
|
|
|
func TestLimiterGroup_Limit(t *testing.T) {
|
|
g := newLimiterGroup(NewLimiter(10), NewLimiter(5), NewLimiter(20))
|
|
if g.Limit() != 5 {
|
|
t.Fatalf("expected 5, got %d", g.Limit())
|
|
}
|
|
}
|
|
|
|
func TestLimiterGroup_Set(t *testing.T) {
|
|
g := newLimiterGroup(NewLimiter(5))
|
|
g.Set(100)
|
|
if g.Limit() != 5 {
|
|
t.Fatal("Set should be a no-op on limiterGroup")
|
|
}
|
|
}
|
|
|
|
func TestLimiterGroup_Empty(t *testing.T) {
|
|
g := newLimiterGroup()
|
|
if g.Limit() != 0 {
|
|
t.Fatalf("empty group limit should be 0, got %d", g.Limit())
|
|
}
|
|
n := g.Wait(context.Background(), 100)
|
|
if n != 100 {
|
|
t.Fatalf("empty group Wait should return n unchanged, got %d", n)
|
|
}
|
|
}
|
|
|
|
func TestLimiterGroup_String(t *testing.T) {
|
|
g := newLimiterGroup(NewLimiter(5))
|
|
if g.String() != "[5]" {
|
|
t.Fatalf("expected '[5]', got '%s'", g.String())
|
|
}
|
|
}
|
|
|
|
var (
|
|
_ limiter.Limiter = (*llimiter)(nil)
|
|
_ limiter.Limiter = (*limiterGroup)(nil)
|
|
)
|