4566704
9ff1a0e816
- 包子句 IPRateLimiter → ipratelimiter(符合 Go 小写命名惯例), 消费方 限定符同步; 新增依赖 golang.org/x/time - 附测试、examples/ipratelimiter 例程与包 README
60 lines
1.3 KiB
Go
60 lines
1.3 KiB
Go
package ipratelimiter
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
func TestNewIPRateLimiter(t *testing.T) {
|
|
l := NewIPRateLimiter(rate.Limit(100), 100)
|
|
if l == nil {
|
|
t.Fatal("NewIPRateLimiter 返回 nil")
|
|
}
|
|
lim := l.GetLimiter("10.1.1.1")
|
|
if lim == nil {
|
|
t.Fatal("GetLimiter 不应返回 nil")
|
|
}
|
|
}
|
|
|
|
func TestGetLimiterSameInstance(t *testing.T) {
|
|
l := NewIPRateLimiter(rate.Limit(100), 100)
|
|
|
|
lim1 := l.GetLimiter("192.168.1.1")
|
|
lim2 := l.GetLimiter("192.168.1.1")
|
|
if lim1 != lim2 {
|
|
t.Error("同一 IP 的 GetLimiter 应返回相同实例")
|
|
}
|
|
}
|
|
|
|
func TestGetLimiterDifferentIP(t *testing.T) {
|
|
l := NewIPRateLimiter(rate.Limit(100), 100)
|
|
|
|
lim1 := l.GetLimiter("192.168.1.1")
|
|
lim2 := l.GetLimiter("10.0.0.1")
|
|
if lim1 == lim2 {
|
|
t.Error("不同 IP 的 GetLimiter 应返回不同实例")
|
|
}
|
|
}
|
|
|
|
func TestAddIP(t *testing.T) {
|
|
l := NewIPRateLimiter(rate.Limit(100), 100)
|
|
|
|
lim1 := l.AddIP("172.16.0.1")
|
|
lim2 := l.GetLimiter("172.16.0.1")
|
|
if lim1 != lim2 {
|
|
t.Error("AddIP 后 GetLimiter 应返回同一实例")
|
|
}
|
|
}
|
|
|
|
func TestAllow(t *testing.T) {
|
|
// 大速率：短时间内必然放行
|
|
l := NewIPRateLimiter(rate.Limit(1000), 1000)
|
|
lim := l.GetLimiter("192.168.1.100")
|
|
for i := 0; i < 10; i++ {
|
|
if !lim.Allow() {
|
|
t.Fatalf("第 %d 次请求被拒绝", i+1)
|
|
}
|
|
}
|
|
}
|