eh-downloader/internal/eh/limiter_test.go

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)
	}
}