167 lines
4.8 KiB
Go
167 lines
4.8 KiB
Go
package hd4k
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// 全局冷却：设置后所有 Wait 都要等，冷却过后才放行。
|
|
//
|
|
// 这是 429 的应对核心 —— 各 worker 独立退避会持续往已限流的服务器打请求，
|
|
// 一起停下才缓得过来。
|
|
func TestLimiterCooldownBlocksAllWaiters(t *testing.T) {
|
|
l := NewLimiter(0) // 不限流，只有冷却起作用
|
|
|
|
l.Cooldown(250 * time.Millisecond)
|
|
if left := l.CooldownLeft(); left <= 0 || left > 250*time.Millisecond {
|
|
t.Fatalf("CooldownLeft = %v, 应在 (0, 250ms] 之间", left)
|
|
}
|
|
|
|
start := time.Now()
|
|
// 三个并发等待者都应被挡住
|
|
for i := 0; i < 3; i++ {
|
|
if err := l.Wait(context.Background()); err != nil {
|
|
t.Fatalf("Wait 返回错误: %v", err)
|
|
}
|
|
}
|
|
// 因为是串行调用，实际只验证最后一个放行时刻；
|
|
// 关键是「不能立刻放行」
|
|
if d := time.Since(start); d < 200*time.Millisecond {
|
|
t.Errorf("冷却没生效，只等了 %v 就放行", d)
|
|
}
|
|
}
|
|
|
|
// 冷却期间必须能被 context 取消（点「停止」不该继续卡着等）
|
|
func TestLimiterCooldownRespectsCancel(t *testing.T) {
|
|
l := NewLimiter(0)
|
|
l.Cooldown(5 * time.Second)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
go func() {
|
|
time.Sleep(80 * time.Millisecond)
|
|
cancel()
|
|
}()
|
|
|
|
start := time.Now()
|
|
err := l.Wait(ctx)
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("应返回 context.Canceled, 实际 %v", err)
|
|
}
|
|
if d := time.Since(start); d > 2*time.Second {
|
|
t.Errorf("取消后还等了 %v，取消没生效", d)
|
|
}
|
|
}
|
|
|
|
// 冷却到期后必须真的放行（不能把自己永久锁死）
|
|
func TestLimiterCooldownExpires(t *testing.T) {
|
|
l := NewLimiter(0)
|
|
l.Cooldown(100 * time.Millisecond)
|
|
time.Sleep(150 * time.Millisecond)
|
|
|
|
if left := l.CooldownLeft(); left != 0 {
|
|
t.Errorf("冷却过期后 CooldownLeft = %v, 想要 0", left)
|
|
}
|
|
start := time.Now()
|
|
if err := l.Wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if d := time.Since(start); d > 100*time.Millisecond {
|
|
t.Errorf("冷却已过期却还等了 %v", d)
|
|
}
|
|
}
|
|
|
|
// 多次触发取最晚，而不是被短的覆盖
|
|
func TestLimiterCooldownTakesMax(t *testing.T) {
|
|
l := NewLimiter(0)
|
|
l.Cooldown(3 * time.Second)
|
|
l.Cooldown(300 * time.Millisecond) // 更短，不该覆盖
|
|
if left := l.CooldownLeft(); left < 2*time.Second {
|
|
t.Errorf("短冷却覆盖了长冷却, CooldownLeft = %v", left)
|
|
}
|
|
}
|
|
|
|
// interval=0（默认）时冷却仍必须生效 ——
|
|
// 否则 Cooldown 在默认配置下就是死代码
|
|
func TestLimiterCooldownWorksWithZeroInterval(t *testing.T) {
|
|
l := NewLimiter(0)
|
|
l.Cooldown(200 * time.Millisecond)
|
|
start := time.Now()
|
|
if err := l.Wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if d := time.Since(start); d < 150*time.Millisecond {
|
|
t.Errorf("interval=0 时冷却被跳过了，只等了 %v", d)
|
|
}
|
|
}
|
|
|
|
func TestParseRetryAfter(t *testing.T) {
|
|
cases := []struct {
|
|
in string
|
|
want time.Duration
|
|
}{
|
|
{"", 0},
|
|
{"0", 0},
|
|
{"-5", 0},
|
|
{"15", 15 * time.Second},
|
|
{"120", 2 * time.Minute},
|
|
{"99999", time.Hour}, // 超 1 小时封顶
|
|
// HTTP 日期格式
|
|
{time.Now().Add(30 * time.Second).UTC().Format("Mon, 02 Jan 2006 15:04:05 GMT"), 30 * time.Second},
|
|
// 已过期的日期
|
|
{time.Now().Add(-time.Minute).UTC().Format("Mon, 02 Jan 2006 15:04:05 GMT"), 0},
|
|
{"nonsense", 0},
|
|
}
|
|
for _, c := range cases {
|
|
got := parseRetryAfter(c.in)
|
|
// 日期解析有时序误差，放宽 2 秒
|
|
diff := got - c.want
|
|
if diff < 0 {
|
|
diff = -diff
|
|
}
|
|
if diff > 2*time.Second {
|
|
t.Errorf("parseRetryAfter(%q) = %v, 想要 %v", c.in, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestIsRateLimited(t *testing.T) {
|
|
for _, code := range []int{429, 503, 403} {
|
|
err := &StatusError{Code: code}
|
|
if !IsRateLimited(err) {
|
|
t.Errorf("HTTP %d 应判定为限流", code)
|
|
}
|
|
}
|
|
for _, code := range []int{404, 500, 200} {
|
|
err := &StatusError{Code: code}
|
|
if IsRateLimited(err) {
|
|
t.Errorf("HTTP %d 不该判定为限流", code)
|
|
}
|
|
}
|
|
if IsRateLimited(nil) {
|
|
t.Error("nil 不该是限流")
|
|
}
|
|
if IsRateLimited(errors.New("connection reset")) {
|
|
t.Error("普通错误不该是限流")
|
|
}
|
|
// 错误链要能穿透（fmt.Errorf %w 包一层也得认出来）
|
|
wrapped := errors.New("outer: " + (&StatusError{Code: 429}).Error())
|
|
// 注意：字符串匹配不算数，必须是 *StatusError 才认
|
|
if IsRateLimited(wrapped) {
|
|
t.Error("只有字符串特征的错误不该被判定为限流（需保留类型）")
|
|
}
|
|
}
|
|
|
|
func TestStatusErrorMessage(t *testing.T) {
|
|
e := &StatusError{Code: 429, RetryAfter: 15 * time.Second, URL: "https://x/y.webp"}
|
|
if got := e.Error(); got == "" {
|
|
t.Error("错误信息不能为空")
|
|
}
|
|
// 不带 RetryAfter 时不该出现「要求等」字样
|
|
e2 := &StatusError{Code: 500, URL: "https://x/y.webp"}
|
|
if got := e2.Error(); got == "" {
|
|
t.Error("错误信息不能为空")
|
|
}
|
|
}
|