129 lines
3.6 KiB
Go
129 lines
3.6 KiB
Go
package fetch
|
||
|
||
import (
|
||
"context"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"sync"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
)
|
||
|
||
func TestGetSuccess(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
_, _ = w.Write([]byte("ok"))
|
||
}))
|
||
defer srv.Close()
|
||
|
||
c := New(Options{MaxConcurrency: 1, Retries: 0})
|
||
defer c.Close()
|
||
|
||
body, code, err := c.Get(context.Background(), srv.URL)
|
||
if err != nil || code != http.StatusOK || string(body) != "ok" {
|
||
t.Fatalf("期望 200/ok，得到 code=%d body=%q err=%v", code, body, err)
|
||
}
|
||
}
|
||
|
||
// TestGetRetriesOn5xx 5xx 应重试到上限后失败（首次 + 2 次重试 = 3 次请求）。
|
||
func TestGetRetriesOn5xx(t *testing.T) {
|
||
var calls int32
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
atomic.AddInt32(&calls, 1)
|
||
w.WriteHeader(http.StatusInternalServerError)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
c := New(Options{MaxConcurrency: 1, Retries: 2, Backoff: time.Millisecond})
|
||
defer c.Close()
|
||
|
||
_, code, err := c.Get(context.Background(), srv.URL)
|
||
if err == nil {
|
||
t.Fatal("500 最终应返回 error")
|
||
}
|
||
if code != http.StatusInternalServerError {
|
||
t.Fatalf("最终状态码应为 500，得到 %d", code)
|
||
}
|
||
if got := atomic.LoadInt32(&calls); got != 3 {
|
||
t.Fatalf("500 应请求 3 次（首次+2 重试），实际 %d", got)
|
||
}
|
||
}
|
||
|
||
// TestGetNoRetryOn4xx 4xx（非 429）不应重试。
|
||
func TestGetNoRetryOn4xx(t *testing.T) {
|
||
var calls int32
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
atomic.AddInt32(&calls, 1)
|
||
w.WriteHeader(http.StatusNotFound)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
c := New(Options{MaxConcurrency: 1, Retries: 3, Backoff: time.Millisecond})
|
||
defer c.Close()
|
||
|
||
_, code, err := c.Get(context.Background(), srv.URL)
|
||
if err == nil || code != http.StatusNotFound {
|
||
t.Fatalf("404 应报错且状态码 404，得到 code=%d err=%v", code, err)
|
||
}
|
||
if got := atomic.LoadInt32(&calls); got != 1 {
|
||
t.Fatalf("404 不应重试（只请求 1 次），实际 %d", got)
|
||
}
|
||
}
|
||
|
||
// TestGetRetriesOn429 429（限流）应重试，后续成功则返回成功。
|
||
func TestGetRetriesOn429(t *testing.T) {
|
||
var calls int32
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
if atomic.AddInt32(&calls, 1) == 1 {
|
||
w.WriteHeader(http.StatusTooManyRequests)
|
||
return
|
||
}
|
||
_, _ = w.Write([]byte("ok"))
|
||
}))
|
||
defer srv.Close()
|
||
|
||
c := New(Options{MaxConcurrency: 1, Retries: 2, Backoff: time.Millisecond})
|
||
defer c.Close()
|
||
|
||
body, code, err := c.Get(context.Background(), srv.URL)
|
||
if err != nil || code != http.StatusOK || string(body) != "ok" {
|
||
t.Fatalf("429 后重试应成功，得到 code=%d body=%q err=%v", code, body, err)
|
||
}
|
||
}
|
||
|
||
// TestConcurrencyLimit 并发上限不得被突破（修复 MaxCo² 回归的守门测试）。
|
||
func TestConcurrencyLimit(t *testing.T) {
|
||
var inflight, peak int32
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
n := atomic.AddInt32(&inflight, 1)
|
||
for {
|
||
old := atomic.LoadInt32(&peak)
|
||
if n <= old || atomic.CompareAndSwapInt32(&peak, old, n) {
|
||
break
|
||
}
|
||
}
|
||
time.Sleep(20 * time.Millisecond)
|
||
atomic.AddInt32(&inflight, -1)
|
||
_, _ = w.Write([]byte("ok"))
|
||
}))
|
||
defer srv.Close()
|
||
|
||
const limit = 2
|
||
c := New(Options{MaxConcurrency: limit, Retries: 0})
|
||
defer c.Close()
|
||
|
||
var wg sync.WaitGroup
|
||
for i := 0; i < 8; i++ {
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
_, _, _ = c.Get(context.Background(), srv.URL)
|
||
}()
|
||
}
|
||
wg.Wait()
|
||
|
||
if m := atomic.LoadInt32(&peak); m > limit {
|
||
t.Fatalf("并发峰值不应超过 %d，实际 %d", limit, m)
|
||
}
|
||
}
|