package eh // 限流器测试:验证「所有请求之间至少间隔 N 毫秒」, // 以及并发情况下也不会穿透(这正是对齐参考脚本的关键行为)。 import ( "context" "net/http" "net/http/httptest" "sync" "testing" "time" ) // testWriter 是 httptest 响应写的最小接口 type testWriter interface { WriteHeader(int) Write([]byte) (int, error) } func newTestServer(t *testing.T, h func(testWriter)) *httptest.Server { t.Helper() return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { h(w) })) } // interval=0 必须完全旁路,不能有任何额外延迟 func TestLimiterZeroIsNoop(t *testing.T) { l := NewLimiter(0) start := time.Now() var wg sync.WaitGroup for i := 0; i < 50; i++ { wg.Add(1) go func() { defer wg.Done() l.Wait(context.Background()) }() } wg.Wait() if elapsed := time.Since(start); elapsed > 200*time.Millisecond { t.Errorf("interval=0 时不该有等待,实际耗时 %v", elapsed) } } // 串行请求:每次都该等满间隔 func TestLimiterSerialSpacing(t *testing.T) { const interval = 100 * time.Millisecond l := NewLimiter(interval) ctx := context.Background() start := time.Now() for i := 0; i < 3; i++ { if err := l.Wait(ctx); err != nil { t.Fatal(err) } } elapsed := time.Since(start) // 3 个请求 => 等 2 次间隔(第 1 个立即放行) wantMin := 2 * interval if elapsed < wantMin { t.Errorf("3 个请求应至少等 %v,实际 %v(限流没生效)", wantMin, elapsed) } } // 并发请求也不能穿透:N 个同时来,应该被拉成串行 func TestLimiterConcurrentSpacing(t *testing.T) { const ( interval = 50 * time.Millisecond n = 5 ) l := NewLimiter(interval) ctx := context.Background() start := time.Now() release := make([]time.Time, n) var wg sync.WaitGroup for i := 0; i < n; i++ { wg.Add(1) go func(idx int) { defer wg.Done() if err := l.Wait(ctx); err != nil { t.Error(err) return } release[idx] = time.Now() }(i) } wg.Wait() elapsed := time.Since(start) // n 个请求 => 至少等 (n-1) 个间隔 wantMin := time.Duration(n-1) * interval if elapsed < wantMin { t.Errorf("%d 个并发请求应至少耗时 %v,实际 %v(限流被穿透)", n, wantMin, elapsed) } } // 取消要能立即返回,不能卡在排队里 func TestLimiterCancel(t *testing.T) { l := NewLimiter(500 * time.Millisecond) ctx, cancel := context.WithCancel(context.Background()) // 先占掉一个名额,让后面的排在后面 l.Wait(context.Background()) go func() { time.Sleep(50 * time.Millisecond) cancel() }() start := time.Now() err := l.Wait(ctx) if err == nil { t.Error("取消后 Wait 应该返回错误") } if elapsed := time.Since(start); elapsed > 200*time.Millisecond { t.Errorf("取消后应立即返回,实际等��� %v", elapsed) } } // 运行时改间隔要生效 func TestLimiterSetInterval(t *testing.T) { l := NewLimiter(0) if l.Interval() != 0 { t.Error("初始间隔应为 0") } l.SetInterval(200 * time.Millisecond) if l.Interval() != 200*time.Millisecond { t.Errorf("改间隔失败: %v", l.Interval()) } // 负数要被夹成 0 l.SetInterval(-100 * time.Millisecond) if l.Interval() != 0 { t.Errorf("负间隔应夹成 0,实际 %v", l.Interval()) } } // 负数构造要夹成 0 func TestNewLimiterNegative(t *testing.T) { l := NewLimiter(-time.Second) if l.Interval() != 0 { t.Errorf("负间隔应夹成 0,实际 %v", l.Interval()) } } // Client 的限流开关 func TestClientRequestInterval(t *testing.T) { c := NewClient(5*time.Second, "", "off", "auto") if c.RequestInterval() != 0 { t.Error("默认不该限流") } c.SetRequestInterval(150 * time.Millisecond) if c.RequestInterval() != 150*time.Millisecond { t.Errorf("设置失败: %v", c.RequestInterval()) } c.SetRequestInterval(-5) if c.RequestInterval() != 0 { t.Errorf("负数应夹成 0,实际 %v", c.RequestInterval()) } } // 限流要真的作用到 HTTP 请求上(起个服务器实测) func TestLimiterAppliesToRequests(t *testing.T) { srv := newTestServer(t, func(w testWriter) { w.WriteHeader(200) w.Write([]byte("ok")) }) defer srv.Close() const interval = 60 * time.Millisecond c := NewClient(5*time.Second, "", "off", "auto") c.SetRequestInterval(interval) start := time.Now() for i := 0; i < 4; i++ { if _, err := c.getText(context.Background(), srv.URL+"/x"); err != nil { t.Fatalf("第 %d 次请求失败: %v", i+1, err) } } elapsed := time.Since(start) // 4 个请求 = 8 次 HTTP 调用?不 —— 这里只有 getText,一次一个请求, // 但下载会额外算一次。总共 4 次请求,至少 3 个间隔。 wantMin := 3 * interval if elapsed < wantMin { t.Errorf("4 次请求应至少耗时 %v,实际 %v", wantMin, elapsed) } }