79 lines
1.8 KiB
Go
79 lines
1.8 KiB
Go
package batchfetch
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestFetchAll(t *testing.T) {
|
|
var cur, peak int64
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
n := atomic.AddInt64(&cur, 1)
|
|
defer atomic.AddInt64(&cur, -1)
|
|
for {
|
|
p := atomic.LoadInt64(&peak)
|
|
if n <= p || atomic.CompareAndSwapInt64(&peak, p, n) {
|
|
break
|
|
}
|
|
}
|
|
id := strings.TrimPrefix(r.URL.Path, "/items/")
|
|
switch id {
|
|
case "slow":
|
|
select {
|
|
case <-time.After(time.Second):
|
|
case <-r.Context().Done():
|
|
}
|
|
case "bad":
|
|
http.Error(w, "boom", http.StatusInternalServerError)
|
|
return
|
|
default:
|
|
time.Sleep(20 * time.Millisecond)
|
|
}
|
|
_, _ = w.Write([]byte("ok:" + id))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
ids := []string{"a", "b", "slow", "c", "bad", "d", "e", "f"}
|
|
res := FetchAll(context.Background(), srv.Client(), srv.URL, ids, 3, 200*time.Millisecond)
|
|
|
|
if peak > 3 {
|
|
t.Fatalf("并发数 %d 超过上限 3", peak)
|
|
}
|
|
for i, r := range res {
|
|
if r.ID != ids[i] {
|
|
t.Fatalf("结果顺序不对：%v", r.ID)
|
|
}
|
|
switch r.ID {
|
|
case "slow":
|
|
if !errors.Is(r.Err, context.DeadlineExceeded) {
|
|
t.Fatalf("slow 应该超时，得到 %v", r.Err)
|
|
}
|
|
case "bad":
|
|
if r.Err == nil || !strings.Contains(r.Err.Error(), "500") {
|
|
t.Fatalf("bad 应该返回 500 错误，得到 %v", r.Err)
|
|
}
|
|
default:
|
|
if r.Err != nil || r.Body != "ok:"+r.ID {
|
|
t.Fatalf("%s: %v %q", r.ID, r.Err, r.Body)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestCancelledContext(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
res := FetchAll(ctx, http.DefaultClient, "http://127.0.0.1:1", []string{"x", "y"}, 1, time.Second)
|
|
for _, r := range res {
|
|
if r.Err == nil {
|
|
t.Fatal("已取消的 context 应该让所有请求失败")
|
|
}
|
|
}
|
|
}
|