scatter_test.go
Rama master · réplica de lectura del repositorio original.
package aastro
import (
"bytes"
"context"
"io"
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/go-chi/chi/v5"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/voidrunner3074/aastro/internal/circuitbreaker"
)
func TestScatter_DispatchingToMultipleUpstreams(t *testing.T) {
t.Run("returns responses from all of them", func(t *testing.T) {
serverA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.Copy(io.Discard, r.Body)
_, _ = w.Write([]byte("A"))
}))
defer serverA.Close()
serverB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.Copy(io.Discard, r.Body)
_, _ = w.Write([]byte("B"))
}))
defer serverB.Close()
f := newTestFlow([]upstream{
newTestUpstream(serverA.URL),
newTestUpstream(serverB.URL),
})
results := newTestScatter().scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
require.Len(t, results, 2)
assert.Nil(t, results[0].err)
assert.Nil(t, results[1].err)
assert.Equal(t, "A", string(results[0].body))
assert.Equal(t, "B", string(results[1].body))
})
}
func TestScatter_Call(t *testing.T) {
t.Run("calls the single upstream directly, without a result slice", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte("solo"))
}))
defer server.Close()
f := newTestFlow([]upstream{newTestUpstream(server.URL)})
resp, ok := newTestScatter().call(f, httptest.NewRequest(http.MethodGet, "/", nil))
assert.True(t, ok)
assert.Nil(t, resp.err)
assert.Equal(t, "solo", string(resp.body))
})
t.Run("returns ok=false for an oversized request body", func(t *testing.T) {
f := newTestFlow([]upstream{newTestUpstream("http://localhost")})
oversizedBody := bytes.Repeat([]byte("x"), maxBodySize+1)
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(oversizedBody))
_, ok := newTestScatter().call(f, req)
assert.False(t, ok)
})
}
func TestScatter_ForwardingRequestData(t *testing.T) {
t.Run("forwards POST body to upstream", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, _ := io.ReadAll(r.Body)
_, _ = w.Write(body)
}))
defer server.Close()
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withMethod(http.MethodPost)),
})
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString("hello"))
results := newTestScatter().scatter(f, req)
require.Len(t, results, 1)
assert.Nil(t, results[0].err)
assert.Equal(t, "hello", string(results[0].body))
})
t.Run("returns nil for an oversized request body", func(t *testing.T) {
f := newTestFlow([]upstream{
newTestUpstream("http://localhost"),
})
oversizedBody := bytes.Repeat([]byte("x"), maxBodySize+1)
req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(oversizedBody))
assert.Nil(t, newTestScatter().scatter(f, req))
})
t.Run("forwards configured query and header values", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
q := r.URL.Query().Get("foo")
h := r.Header.Get("X-Test")
_, _ = w.Write([]byte(q + "-" + h))
}))
defer server.Close()
f := newTestFlow([]upstream{
newTestUpstream(server.URL,
withForwardQueries("foo"),
withForwardHeaders("X-Test"),
),
})
req := httptest.NewRequest(http.MethodGet, "/?foo=bar", nil)
req.Header.Set("X-Test", "baz")
results := newTestScatter().scatter(f, req)
require.Len(t, results, 1)
assert.Equal(t, "bar-baz", string(results[0].body))
})
t.Run("forwards path params as query when configured", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(r.URL.Query().Get("user_id")))
}))
defer server.Close()
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withForwardParams("user_id")),
})
rctx := chi.NewRouteContext()
rctx.URLParams.Add("user_id", "42")
req := httptest.NewRequest(http.MethodGet, "/users/42", nil)
req = req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, rctx))
results := newTestScatter().scatter(f, req)
require.Len(t, results, 1)
assert.Equal(t, "42", string(results[0].body))
})
t.Run("expands path params in upstream path template", func(t *testing.T) {
var receivedPath string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
receivedPath = r.URL.Path
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withPath("/orders/{order_id}")),
})
rctx := chi.NewRouteContext()
rctx.URLParams.Add("order_id", "99")
req := httptest.NewRequest(http.MethodGet, "/orders/99", nil)
req = req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, rctx))
newTestScatter().scatter(f, req)
assert.Equal(t, "/orders/99", receivedPath)
})
}
func TestScatter_UpstreamPolicies(t *testing.T) {
t.Run("rejects empty body when require_body is set", func(t *testing.T) {
serverWithBody := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte(`{"ok":true}`))
}))
defer serverWithBody.Close()
serverNoBody := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusNoContent)
}))
defer serverNoBody.Close()
f := newTestFlow([]upstream{
newTestUpstream(serverWithBody.URL, withPolicy(upstreamPolicy{requireBody: true})),
newTestUpstream(serverNoBody.URL, withPolicy(upstreamPolicy{requireBody: true})),
})
results := newTestScatter().scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
require.Len(t, results, 2)
assert.Nil(t, results[0].err)
require.NotNil(t, results[1].err)
assert.EqualError(t, results[1].err.Unwrap(), "empty body not allowed by upstream policy")
})
t.Run("rejects responses larger than max_response_body_size", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte("abcdefghijklmnopqrstuvwxyz"))
}))
defer server.Close()
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withPolicy(upstreamPolicy{maxResponseBodySize: 10})),
})
results := newTestScatter().scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
require.Len(t, results, 1)
require.NotNil(t, results[0].err)
assert.Equal(t, upstreamBodyTooLarge, results[0].err.kind)
})
t.Run("classifies upstream timeouts", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
time.Sleep(600 * time.Millisecond)
}))
defer server.Close()
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withTimeout(100*time.Millisecond)),
})
results := newTestScatter().scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
require.Len(t, results, 1)
require.NotNil(t, results[0].err)
assert.Equal(t, upstreamTimeout, results[0].err.kind)
})
}
func TestScatter_RetryPolicy(t *testing.T) {
t.Run("succeeds after a few transient failures", func(t *testing.T) {
var attempts atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
if attempts.Add(1) <= 2 {
w.WriteHeader(http.StatusInternalServerError)
return
}
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withPolicy(upstreamPolicy{
retry: retryPolicy{
maxRetries: 3,
retryOnStatuses: []int{http.StatusInternalServerError},
backoffDelay: 10 * time.Millisecond,
},
})),
})
results := newTestScatter().scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
require.Len(t, results, 1)
assert.Nil(t, results[0].err)
assert.Equal(t, int32(3), attempts.Load())
})
t.Run("still retries when the upstream has no explicit method (falls back to the request's)", func(t *testing.T) {
// Regression test: shouldRetry must judge idempotency against the
// effective method (config method, falling back to the original
// request's), not the raw config value - an unset upstream method
// used to make isIdempotent("") always false, silently disabling
// retries for the common case of not overriding the method.
var attempts atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
if attempts.Add(1) <= 2 {
w.WriteHeader(http.StatusInternalServerError)
return
}
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withMethod(""), withPolicy(upstreamPolicy{
retry: retryPolicy{
maxRetries: 3,
retryOnStatuses: []int{http.StatusInternalServerError},
backoffDelay: 10 * time.Millisecond,
},
})),
})
results := newTestScatter().scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
require.Len(t, results, 1)
assert.Nil(t, results[0].err)
assert.Equal(t, int32(3), attempts.Load())
})
t.Run("gives up after exhausting max retries", func(t *testing.T) {
var attempts atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
attempts.Add(1)
w.WriteHeader(http.StatusInternalServerError)
}))
defer server.Close()
maxRetries := 3
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withPolicy(upstreamPolicy{
retry: retryPolicy{
maxRetries: maxRetries,
retryOnStatuses: []int{http.StatusInternalServerError},
backoffDelay: 10 * time.Millisecond,
},
})),
})
results := newTestScatter().scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
require.Len(t, results, 1)
require.NotNil(t, results[0].err)
assert.Equal(t, upstreamBadStatus, results[0].err.kind)
assert.Equal(t, int32(maxRetries+1), attempts.Load())
})
}
func TestScatter_CircuitBreaker(t *testing.T) {
t.Run("opens after configured failure count and rejects further requests", func(t *testing.T) {
var upstreamCalls atomic.Int32
maxFailures := 3
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
upstreamCalls.Add(1)
w.WriteHeader(http.StatusInternalServerError)
}))
defer server.Close()
cb := circuitbreaker.New(maxFailures, 100*time.Millisecond)
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withCircuitBreaker(cb)),
})
d := newTestScatter()
results := make([]upstreamResponse, 5)
for i := range results {
responses := d.scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
results[i] = responses[0]
}
for i := range maxFailures {
require.NotNil(t, results[i].err)
assert.Equal(t, upstreamBadStatus, results[i].err.kind)
}
for i := maxFailures; i < 5; i++ {
require.NotNil(t, results[i].err)
assert.Equal(t, upstreamCircuitOpen, results[i].err.kind)
}
assert.Equal(t, int32(maxFailures), upstreamCalls.Load())
})
t.Run("closes after reset timeout and serves successful requests", func(t *testing.T) {
resetTimeout := 100 * time.Millisecond
cb := circuitbreaker.New(1, resetTimeout)
var calls atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
if calls.Add(1) == 1 {
w.WriteHeader(http.StatusInternalServerError)
return
}
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
f := newTestFlow([]upstream{
newTestUpstream(server.URL, withCircuitBreaker(cb)),
})
d := newTestScatter()
r1 := d.scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
require.NotNil(t, r1[0].err)
assert.Equal(t, upstreamBadStatus, r1[0].err.kind)
r2 := d.scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
require.NotNil(t, r2[0].err)
assert.Equal(t, upstreamCircuitOpen, r2[0].err.kind)
time.Sleep(resetTimeout + 20*time.Millisecond)
r3 := d.scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
assert.Nil(t, r3[0].err)
})
}
func TestScatter_LoadBalancing(t *testing.T) {
t.Run("distributes round-robin evenly across hosts", func(t *testing.T) {
var callsA, callsB atomic.Int32
serverA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
callsA.Add(1)
}))
defer serverA.Close()
serverB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
callsB.Add(1)
}))
defer serverB.Close()
f := newTestFlow([]upstream{
newTestUpstream("",
withHosts(serverA.URL, serverB.URL),
withLBMode(lbModeRoundRobin, 2),
),
})
d := newTestScatter()
for range 4 {
d.scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
}
assert.Equal(t, int32(2), callsA.Load())
assert.Equal(t, int32(2), callsB.Load())
})
t.Run("least-connections sends more traffic to faster host", func(t *testing.T) {
var callsA, callsB atomic.Int32
serverA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
callsA.Add(1)
time.Sleep(100 * time.Millisecond)
}))
defer serverA.Close()
serverB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
callsB.Add(1)
}))
defer serverB.Close()
f := newTestFlow([]upstream{
newTestUpstream("",
withHosts(serverA.URL, serverB.URL),
withLBMode(lbModeLeastConns, 2),
),
})
d := newTestScatter()
var wg sync.WaitGroup
for range 10 {
wg.Add(1)
go func() {
defer wg.Done()
d.scatter(f, httptest.NewRequest(http.MethodGet, "/", nil))
}()
time.Sleep(20 * time.Millisecond)
}
wg.Wait()
assert.Greater(t, callsB.Load(), callsA.Load())
})
}