129 lines
3.5 KiB
Go
129 lines
3.5 KiB
Go
package steamcache
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestBandwidthFairShareRates(t *testing.T) {
|
|
b := newClientBandwidthLimiter(1000, 0)
|
|
lim1, rel1 := b.acquire("1.1.1.1")
|
|
defer rel1()
|
|
lim2, rel2 := b.acquire("2.2.2.2")
|
|
defer rel2()
|
|
if lim1 == nil || lim2 == nil {
|
|
t.Fatal("expected limiters")
|
|
}
|
|
// With 2 active clients, each should get ~500 bytes/sec.
|
|
got1 := float64(lim1.Limit())
|
|
got2 := float64(lim2.Limit())
|
|
if got1 < 400 || got1 > 600 || got2 < 400 || got2 > 600 {
|
|
t.Fatalf("fair-share rates = %v,%v want ~500", got1, got2)
|
|
}
|
|
rel2()
|
|
// After release, sole client should get full uplink.
|
|
lim1b, rel1b := b.acquire("1.1.1.1")
|
|
defer rel1b()
|
|
if float64(lim1b.Limit()) < 900 {
|
|
t.Fatalf("after release limit=%v want ~1000", lim1b.Limit())
|
|
}
|
|
}
|
|
|
|
func TestBandwidthAbsoluteCap(t *testing.T) {
|
|
b := newClientBandwidthLimiter(0, 250)
|
|
lim, rel := b.acquire("9.9.9.9")
|
|
defer rel()
|
|
if lim == nil {
|
|
t.Fatal("expected limiter")
|
|
}
|
|
if float64(lim.Limit()) != 250 {
|
|
t.Fatalf("limit=%v want 250", lim.Limit())
|
|
}
|
|
}
|
|
|
|
func TestBandwidthDisabled(t *testing.T) {
|
|
b := newClientBandwidthLimiter(0, 0)
|
|
lim, rel := b.acquire("9.9.9.9")
|
|
defer rel()
|
|
if lim != nil {
|
|
t.Fatal("expected nil limiter when disabled")
|
|
}
|
|
}
|
|
|
|
func TestLimitedResponseWriterShapes(t *testing.T) {
|
|
var buf bytes.Buffer
|
|
rec := httptest.NewRecorder()
|
|
// Use a custom writer sink via ResponseRecorder is fine; WaitN will delay.
|
|
lim := newClientBandwidthLimiter(0, 2000) // 2KB/s
|
|
l, rel := lim.acquire("127.0.0.1")
|
|
defer rel()
|
|
w := &limitedResponseWriter{ResponseWriter: rec, lim: l, ctx: context.Background()}
|
|
payload := bytes.Repeat([]byte("x"), 4000)
|
|
start := time.Now()
|
|
n, err := w.Write(payload)
|
|
elapsed := time.Since(start)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if n != len(payload) {
|
|
t.Fatalf("wrote %d want %d", n, len(payload))
|
|
}
|
|
_ = buf
|
|
// 4000 bytes at 2000 B/s should take ~2s (allow slack for CI).
|
|
if elapsed < 1500*time.Millisecond {
|
|
t.Fatalf("elapsed %v too fast for 2KB/s shaping of 4KB", elapsed)
|
|
}
|
|
if elapsed > 8*time.Second {
|
|
t.Fatalf("elapsed %v unexpectedly slow", elapsed)
|
|
}
|
|
}
|
|
|
|
func TestServeHTTPBandwidthCap(t *testing.T) {
|
|
body := bytes.Repeat([]byte("a"), 3000)
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/octet-stream")
|
|
w.Header().Set("Content-Length", "3000")
|
|
w.WriteHeader(200)
|
|
_, _ = w.Write(body)
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
|
|
sc, err := NewWithOptions(Options{
|
|
Address: "127.0.0.1:0",
|
|
MemorySize: "1MB",
|
|
DiskSize: "0",
|
|
Upstream: upstream.URL,
|
|
MemoryGC: "lru",
|
|
DiskGC: "lru",
|
|
MaxConcurrentRequests: 20,
|
|
MaxRequestsPerClient: 10,
|
|
MaxObjectSize: "0",
|
|
MaxBytesPerClientPerSec: 1500, // 1.5KB/s
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewWithOptions: %v", err)
|
|
}
|
|
t.Cleanup(func() { sc.Shutdown() })
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/depot/bw/chunk", nil)
|
|
req.Header.Set("User-Agent", "Valve/Steam HTTP Client 1.0")
|
|
rr := httptest.NewRecorder()
|
|
start := time.Now()
|
|
sc.ServeHTTP(rr, req)
|
|
elapsed := time.Since(start)
|
|
if rr.Code != 200 {
|
|
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
|
}
|
|
got := rr.Body.Bytes()
|
|
if len(got) != len(body) {
|
|
t.Fatalf("body len=%d want %d", len(got), len(body))
|
|
}
|
|
if elapsed < time.Second {
|
|
t.Fatalf("elapsed %v too fast for shaping", elapsed)
|
|
}
|
|
}
|