helpers_test.go
Rama master · réplica de lectura del repositorio original.
package aastro
import (
"context"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"fmt"
"math/big"
"net"
"net/http"
"os"
"testing"
"time"
"github.com/go-chi/chi/v5"
"github.com/goccy/go-json"
"github.com/stretchr/testify/require"
"go.uber.org/zap"
"github.com/voidrunner3074/aastro/internal/circuitbreaker"
"github.com/voidrunner3074/aastro/internal/metric"
"github.com/voidrunner3074/aastro/sdk"
)
var testMetrics *metric.Metrics
func TestMain(m *testing.M) {
tm, err := metric.New()
if err != nil {
_, _ = fmt.Fprintf(os.Stderr, "init test metrics: %v\n", err)
os.Exit(1)
}
testMetrics = tm
os.Exit(m.Run())
}
// ── Upstream stubs ────────────────────────────────────────────────────────────
type mockUpstream struct {
upstreamName string
}
func (s *mockUpstream) name() string { return s.upstreamName }
func (s *mockUpstream) call(_ context.Context, _ *http.Request, _ []byte) *upstreamResponse {
return &upstreamResponse{}
}
type mockProxyUpstream struct {
upstreamName string
proxyFn func(w http.ResponseWriter, req *http.Request) error
}
func (m *mockProxyUpstream) name() string { return m.upstreamName }
func (m *mockProxyUpstream) call(_ context.Context, _ *http.Request, _ []byte) *upstreamResponse {
return &upstreamResponse{}
}
func (m *mockProxyUpstream) proxy(_ context.Context, w http.ResponseWriter, req *http.Request) error {
return m.proxyFn(w, req)
}
func mockUpstreams(names ...string) []upstream {
upstreams := make([]upstream, len(names))
for i, name := range names {
upstreams[i] = &mockUpstream{upstreamName: name}
}
return upstreams
}
// ── Plugin and middleware mocks ───────────────────────────────────────────────
type mockPlugin struct {
name string
typ sdk.PluginType
fn func(sdk.Context)
}
func (m *mockPlugin) Init(_ map[string]interface{}) error { return nil }
func (m *mockPlugin) Info() sdk.PluginInfo {
return sdk.PluginInfo{
Name: m.name,
Description: "Mock plugin",
Version: "v1",
Author: "test",
}
}
func (m *mockPlugin) Type() sdk.PluginType { return m.typ }
func (m *mockPlugin) Execute(ctx sdk.Context) error {
m.fn(ctx)
return nil
}
type mockMiddleware struct{}
func (m *mockMiddleware) Init(_ map[string]interface{}) error { return nil }
func (m *mockMiddleware) Name() string { return "mockmw" }
func (m *mockMiddleware) Handler(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Add("X-Middleware", "ok")
next.ServeHTTP(w, r)
})
}
// ── Scatter mock ──────────────────────────────────────────────────────────────
type mockScatter struct {
results []upstreamResponse
tooLarge bool
}
func (m *mockScatter) scatter(_ *flow, _ *http.Request) []upstreamResponse {
if m.tooLarge {
return nil
}
return m.results
}
func (m *mockScatter) call(_ *flow, _ *http.Request) (upstreamResponse, bool) {
if m.tooLarge {
return upstreamResponse{}, false
}
return m.results[0], true
}
// ── Builders ──────────────────────────────────────────────────────────────────
func newTestUpstreamConfig(port string) UpstreamConfig {
return UpstreamConfig{
Name: "test_service_" + port,
Hosts: AddrList{"test-service:" + port},
Path: "/builder/test",
Method: http.MethodGet,
Timeout: 5 * time.Second,
}
}
func newTestRouter(flows []flow, d scatter, a aggregator) *Router {
r := &Router{
chiRouter: chi.NewMux(),
scatter: d,
aggregator: a,
flows: flows,
log: zap.NewNop(),
metrics: testMetrics,
}
r.registerFlows()
return r
}
func newTestUpstream(host string, opts ...func(*httpUpstream)) *httpUpstream {
u := &httpUpstream{
cfg: upstreamConfig{
hosts: []string{host},
method: http.MethodGet,
timeout: 500 * time.Millisecond,
},
state: upstreamState{},
metrics: testMetrics,
log: zap.NewNop(),
client: http.DefaultClient,
}
for _, opt := range opts {
opt(u)
}
return u
}
func withMethod(method string) func(*httpUpstream) {
return func(u *httpUpstream) { u.cfg.method = method }
}
func withTimeout(d time.Duration) func(*httpUpstream) {
return func(u *httpUpstream) { u.cfg.timeout = d }
}
func withPolicy(p upstreamPolicy) func(*httpUpstream) {
return func(u *httpUpstream) { u.cfg.policy = p }
}
func withForwardQueries(queries ...string) func(*httpUpstream) {
return func(u *httpUpstream) { u.cfg.forwardQueries = queries }
}
func withForwardHeaders(headers ...string) func(*httpUpstream) {
return func(u *httpUpstream) { u.cfg.forwardHeaders = headers }
}
func withForwardParams(params ...string) func(*httpUpstream) {
return func(u *httpUpstream) { u.cfg.forwardParams = params }
}
func withPath(path string) func(*httpUpstream) {
return func(u *httpUpstream) { u.cfg.path = path }
}
func withHosts(hosts ...string) func(*httpUpstream) {
return func(u *httpUpstream) { u.cfg.hosts = hosts }
}
func withLBMode(mode lbMode, hostCount int) func(*httpUpstream) {
return func(u *httpUpstream) {
u.cfg.lbMode = mode
u.state.activeConnections = make([]int64, hostCount)
}
}
func withCircuitBreaker(cb *circuitbreaker.CircuitBreaker) func(*httpUpstream) {
return func(u *httpUpstream) { u.circuitBreaker = cb }
}
func withTrustedProxies(proxies ...*net.IPNet) func(*httpUpstream) {
return func(u *httpUpstream) { u.cfg.trustedProxies = proxies }
}
func withTLS(tlsConfig *tls.Config) func(*httpUpstream) {
return func(u *httpUpstream) {
u.client = &http.Client{
Transport: &http.Transport{TLSClientConfig: tlsConfig},
}
}
}
// ── TLS fixture ───────────────────────────────────────────────────────────────
type tlsFixture struct {
caCert *x509.Certificate
caKey *ecdsa.PrivateKey
}
func newTLSFixture(t *testing.T) *tlsFixture {
t.Helper()
caKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)
template := &x509.Certificate{
SerialNumber: big.NewInt(1),
Subject: pkix.Name{CommonName: "aastro-test-ca"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(time.Hour),
IsCA: true,
KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
BasicConstraintsValid: true,
}
der, err := x509.CreateCertificate(rand.Reader, template, template, &caKey.PublicKey, caKey)
require.NoError(t, err)
caCert, err := x509.ParseCertificate(der)
require.NoError(t, err)
return &tlsFixture{caCert: caCert, caKey: caKey}
}
func (f *tlsFixture) issueServerCert(t *testing.T) tls.Certificate {
t.Helper()
return f.issueCert(t, certParams{
cn: "test-server",
eku: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
dnsNames: []string{"localhost"},
ipAddresses: []net.IP{net.IPv4(127, 0, 0, 1), net.IPv6loopback},
})
}
func (f *tlsFixture) issueClientCert(t *testing.T, cn string) tls.Certificate {
t.Helper()
return f.issueCert(t, certParams{
cn: cn,
eku: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
})
}
func (f *tlsFixture) certPool() *x509.CertPool {
pool := x509.NewCertPool()
pool.AddCert(f.caCert)
return pool
}
type certParams struct {
cn string
eku []x509.ExtKeyUsage
dnsNames []string
ipAddresses []net.IP
}
func (f *tlsFixture) issueCert(t *testing.T, p certParams) tls.Certificate {
t.Helper()
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)
template := &x509.Certificate{
SerialNumber: big.NewInt(time.Now().UnixNano()),
Subject: pkix.Name{CommonName: p.cn},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(time.Hour),
KeyUsage: x509.KeyUsageDigitalSignature,
ExtKeyUsage: p.eku,
DNSNames: p.dnsNames,
IPAddresses: p.ipAddresses,
}
der, err := x509.CreateCertificate(rand.Reader, template, f.caCert, &key.PublicKey, f.caKey)
require.NoError(t, err)
return tls.Certificate{
Certificate: [][]byte{der},
PrivateKey: key,
}
}
// ── Flow / scatter builders ───────────────────────────────────────────────────
func newTestFlow(upstreams []upstream) *flow {
return &flow{upstreams: upstreams}
}
func newTestScatter() *defaultScatter {
return &defaultScatter{
log: zap.NewNop(),
metrics: testMetrics,
}
}
func streamingFlow(path string, u upstream, plugins ...sdk.Plugin) flow {
return flow{
path: path,
method: http.MethodGet,
streaming: true,
upstreams: []upstream{u},
plugins: plugins,
}
}
// ── Response builders ─────────────────────────────────────────────────────────
func okResponse(body string) upstreamResponse {
return upstreamResponse{body: []byte(body)}
}
func errResponse(kind upstreamErrorKind) upstreamResponse {
return upstreamResponse{
err: &upstreamError{
kind: kind,
err: fmt.Errorf("upstream error: %s", kind),
},
}
}
// ── Request builders ──────────────────────────────────────────────────────────
func requestWithClientIP(method, url, remoteAddr, clientIP string) *http.Request {
req, _ := http.NewRequest(method, url, nil)
req.RemoteAddr = remoteAddr
req = req.WithContext(withClientIP(req.Context(), clientIP))
return req
}
func mustParseCIDR(cidr string) *net.IPNet {
_, ipnet, err := net.ParseCIDR(cidr)
if err != nil {
panic(err)
}
return ipnet
}
// ── Assertions ────────────────────────────────────────────────────────────────
func jsonEqual(t *testing.T, expected string, actual []byte) {
t.Helper()
var e, a any
require.NoError(t, json.Unmarshal([]byte(expected), &e))
require.NoError(t, json.Unmarshal(actual, &a))
require.Equal(t, e, a)
}
func decodeProblem(t *testing.T, body []byte) ProblemDetails {
t.Helper()
var problem ProblemDetails
require.NoError(t, json.Unmarshal(body, &problem))
return problem
}
func decodeJSONInto(data []byte, dst any) error {
return json.Unmarshal(data, dst)
}