spb/internetpressure

1package batcher23import (4	"bytes"5	"compress/gzip"6	"context"7	"encoding/json"8	"errors"9	"log/slog"10	"sync"11	"testing"12	"time"1314	"internetpressure.io/probe-agent/internal/client"15	"internetpressure.io/probe-agent/internal/protocol"16	"internetpressure.io/probe-agent/internal/spool"17)1819// fakePoster records batches and can be switched to failing.20type fakePoster struct {21	mu      sync.Mutex22	batches []protocol.Batch23	fail    error24	calls   int25}2627func (p *fakePoster) PostBatch(_ context.Context, gz []byte, _ time.Duration) (*protocol.BatchResponse, error) {28	p.mu.Lock()29	defer p.mu.Unlock()30	p.calls++31	if p.fail != nil {32		return nil, p.fail33	}34	zr, err := gzip.NewReader(bytes.NewReader(gz))35	if err != nil {36		return nil, err37	}38	var b protocol.Batch39	if err := json.NewDecoder(zr).Decode(&b); err != nil {40		return nil, err41	}42	p.batches = append(p.batches, b)43	return &protocol.BatchResponse{Accepted: len(b.Measurements), ConfigVersion: "v1", ServerTime: protocol.FormatTime(time.Now())}, nil44}4546func (p *fakePoster) setFail(err error) { p.mu.Lock(); p.fail = err; p.mu.Unlock() }47func (p *fakePoster) count() int        { p.mu.Lock(); defer p.mu.Unlock(); return len(p.batches) }4849func newBatcher(t *testing.T, p Poster, hooks Hooks) (*Batcher, *spool.Spool) {50	sp, err := spool.Open(t.TempDir(), 0)51	if err != nil {52		t.Fatal(err)53	}54	return New("ca-qc-01", "0.1.0", p, sp, hooks, slog.New(slog.DiscardHandler)), sp55}5657func meas(n int) []protocol.Measurement {58	out := make([]protocol.Measurement, n)59	for i := range out {60		out[i] = protocol.Measurement{TS: protocol.FormatTime(time.Now()), TargetID: "t", Kind: "http", OK: true}61	}62	return out63}6465func TestBatchSizeLimit(t *testing.T) {66	p := &fakePoster{}67	b, _ := newBatcher(t, p, Hooks{})68	b.Configure(10, 500)69	b.Add(meas(1200)...)70	b.AddTraceroute(protocol.Traceroute{TargetID: "t"})71	b.Flush(context.Background())72	if b.Buffered() != 0 {73		t.Fatalf("queue should be drained, has %d", b.Buffered())74	}75	if p.count() != 3 {76		t.Fatalf("1200 measurements should give 3 batches, got %d", p.count())77	}78	for i, batch := range p.batches {79		if len(batch.Measurements) > 500 {80			t.Fatalf("batch %d has %d > max_batch", i, len(batch.Measurements))81		}82		if batch.ProbeID != "ca-qc-01" || batch.AgentVersion != "0.1.0" || batch.SentAt == "" {83			t.Fatalf("envelope: %+v", batch)84		}85	}86	if len(p.batches[0].Measurements) != 500 || len(p.batches[2].Measurements) != 200 || len(p.batches[0].Traceroutes) != 1 {87		t.Fatalf("split: %d %d %d, trs %d", len(p.batches[0].Measurements), len(p.batches[1].Measurements), len(p.batches[2].Measurements), len(p.batches[0].Traceroutes))88	}89}9091func TestHealthOncePerMinute(t *testing.T) {92	p := &fakePoster{}93	calls := 094	b, _ := newBatcher(t, p, Hooks{Health: func() *protocol.Health { calls++; return &protocol.Health{AgentVersion: "0.1.0"} }})95	b.Add(meas(1)...)96	b.Flush(context.Background())97	b.Add(meas(1)...)98	b.Flush(context.Background())99	if calls != 1 || p.batches[0].Health == nil || p.batches[1].Health != nil {100		t.Fatalf("health should be attached once: calls=%d", calls)101	}102	// Empty queue and health not due → no POST at all.103	n := p.calls104	b.Flush(context.Background())105	if p.calls != n {106		t.Fatal("empty flush must not POST")107	}108}109110func TestSpoolOnFailureAndDrain(t *testing.T) {111	p := &fakePoster{}112	failures := 0113	b, sp := newBatcher(t, p, Hooks{OnFlushFailure: func() { failures++ }})114	p.setFail(&client.HTTPError{Status: 503, Body: "down"})115	b.Add(meas(10)...)116	b.Flush(context.Background())117	if n, _ := sp.Stats(); n != 1 || failures != 1 || !b.InBackoff() {118		t.Fatalf("expected 1 spooled file + backoff, got %d files, failures=%d", n, failures)119	}120	// During backoff new batches go straight to the spool without touching the network.121	calls := p.calls122	b.Add(meas(5)...)123	b.Flush(context.Background())124	if p.calls != calls {125		t.Fatal("must not POST during backoff")126	}127	if n, _ := sp.Stats(); n != 2 {128		t.Fatalf("expected 2 spooled files, got %d", n)129	}130	// Server back: force backoff expiry, drain oldest first, then live.131	p.setFail(nil)132	b.nextTry = time.Now().Add(-time.Second)133	b.Add(meas(3)...)134	b.Flush(context.Background())135	if n, _ := sp.Stats(); n != 0 {136		t.Fatalf("spool should be drained, %d left", n)137	}138	if p.count() != 3 || len(p.batches[0].Measurements) != 10 || len(p.batches[1].Measurements) != 5 || len(p.batches[2].Measurements) != 3 {139		t.Fatalf("delivery order wrong: %d batches", p.count())140	}141	if b.InBackoff() {142		t.Fatal("backoff should be cleared after success")143	}144}145146func TestBackoffGrowsAndCaps(t *testing.T) {147	b, _ := newBatcher(t, &fakePoster{}, Hooks{})148	want := []time.Duration{5 * time.Second, 10 * time.Second, 20 * time.Second, 40 * time.Second, 80 * time.Second, 160 * time.Second, 300 * time.Second, 300 * time.Second}149	for i, w := range want {150		b.fail(errors.New("x"))151		if b.backoff != w {152			t.Fatalf("step %d: backoff %v want %v", i, b.backoff, w)153		}154	}155}156157func TestAuthFailureDrops(t *testing.T) {158	p := &fakePoster{}159	b, sp := newBatcher(t, p, Hooks{})160	p.setFail(&client.HTTPError{Status: 401, Body: `{"detail":"bad signature"}`})161	b.Add(meas(2)...)162	b.Flush(context.Background())163	if n, _ := sp.Stats(); n != 0 || b.InBackoff() || b.Buffered() != 0 {164		t.Fatalf("401 must drop without spool/backoff: files=%d backoff=%v buffered=%d", n, b.InBackoff(), b.Buffered())165	}166	// 401 skew → retried once.167	p.setFail(&client.HTTPError{Status: 401, Body: `{"detail":"timestamp skew"}`})168	calls := p.calls169	b.Add(meas(2)...)170	b.Flush(context.Background())171	if p.calls-calls != 2 {172		t.Fatalf("skew should retry exactly once, got %d calls", p.calls-calls)173	}174}175176func TestStopSpoolsRemainder(t *testing.T) {177	p := &fakePoster{}178	b, sp := newBatcher(t, p, Hooks{})179	b.Configure(10, 5)180	p.setFail(errors.New("dial tcp: connection refused"))181	b.Add(meas(12)...)182	b.Stop()183	if b.Buffered() != 0 {184		t.Fatalf("queue not emptied on stop: %d", b.Buffered())185	}186	if n, _ := sp.Stats(); n != 3 {187		t.Fatalf("expected 3 spooled files on stop, got %d", n)188	}189	if p.calls != 1 {190		t.Fatalf("stop should try the network once, tried %d", p.calls)191	}192}193194func TestMaxBatchKicks(t *testing.T) {195	p := &fakePoster{}196	b, _ := newBatcher(t, p, Hooks{})197	b.Configure(3600, 4)198	ctx, cancel := context.WithCancel(context.Background())199	defer cancel()200	go b.Run(ctx)201	b.Add(meas(4)...)202	deadline := time.Now().Add(2 * time.Second)203	for p.count() == 0 && time.Now().Before(deadline) {204		time.Sleep(10 * time.Millisecond)205	}206	if p.count() != 1 {207		t.Fatal("reaching max_batch should flush immediately")208	}209}210