sift/docs/plans/p6/_files/api/batchrun_test.go

package api

// P6 task 05: the runner wired to fetchTarget and the metrics, with no HTTP yet. The helpers
// here (newBatchServer, pageServer, fakeDNS) are shared with task 06's batch_test.go.

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/auth"
	"git.wntrmute.dev/kyle/sift/internal/batch"
	"git.wntrmute.dev/kyle/sift/internal/dcache"
	"git.wntrmute.dev/kyle/sift/internal/fetch"
	"git.wntrmute.dev/kyle/sift/internal/guard"
)

var loopback = []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}

// fakeDNS resolves internal.test to a private address; every other name fails to resolve.
func fakeDNS(_ context.Context, host string) ([]netip.Addr, error) {
	if host == "internal.test" {
		return []netip.Addr{netip.MustParseAddr("10.0.0.5")}, nil
	}
	return nil, fmt.Errorf("no such host %s", host)
}

// newBatchServer builds a Server with the batch runner wired over an in-memory dcache through
// s.NewBatchRunner, exactly as serve does. callers are the names to provision (each gets a key).
func newBatchServer(t *testing.T, callers ...string) (*Server, map[string]string) {
	t.Helper()
	dir := t.TempDir()
	keysPath := filepath.Join(dir, "keys.toml")
	keys := map[string]string{}
	for _, name := range callers {
		out := filepath.Join(dir, name)
		if err := auth.Add(keysPath, name, "hermes-fleet", out); err != nil {
			t.Fatal(err)
		}
		b, err := os.ReadFile(out)
		if err != nil {
			t.Fatal(err)
		}
		keys[name] = strings.TrimSpace(string(b))
	}
	store, err := auth.Load(keysPath)
	if err != nil {
		t.Fatal(err)
	}
	g := &guard.Guard{AllowPrefixes: loopback, LookupIP: fakeDNS}
	f := &fetch.Fetcher{Guard: g, UA: "sift-test"}
	log := slog.New(slog.NewTextHandler(io.Discard, nil))
	s := New("test", g, f, store, log)
	db, err := dcache.Open("", 1<<20, log)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })
	s.Batch = s.NewBatchRunner(db, batch.Config{})
	t.Cleanup(s.Batch.Close) // before db.Close
	return s, keys
}

// pageServer serves a small HTML page at every path; hold, when non-nil, blocks each request
// until it is closed.
func pageServer(t *testing.T, hold chan struct{}) *httptest.Server {
	t.Helper()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if hold != nil {
			select {
			case <-hold:
			case <-r.Context().Done():
				return
			}
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		fmt.Fprintf(w, `<html><head><title>T</title></head><body><article><p>hello batch %s</p></article></body></html>`, r.URL.Path)
	}))
	t.Cleanup(ts.Close)
	return ts
}

func waitJob(t *testing.T, r *batch.Runner, id string) batch.Job {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		j, _, ok := r.Status(id)
		if ok && (j.Status == "done" || j.Status == "partial" || j.Status == "failed" || j.Status == "cancelled") {
			return j
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatal("job never finished")
	return batch.Job{}
}

func promText(t *testing.T, s *Server) string {
	t.Helper()
	req := httptest.NewRequest(http.MethodGet, "/metrics", nil)
	req.RemoteAddr, req.Host = "127.0.0.1:9", "127.0.0.1:8890"
	rec := httptest.NewRecorder()
	s.ServeHTTP(rec, req)
	return rec.Body.String()
}

// The runner fetches through fetchTarget as the submitting caller; a name that resolves to a
// private address is a per-target "blocked" failure; the hooks reach the metrics.
func TestBatchRunnerWired(t *testing.T) {
	ts := pageServer(t, nil)
	s, _ := newBatchServer(t, "test")
	c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
	id, reused, err := s.Batch.Submit(context.Background(), c, dcache.ClassFleet, []batch.Target{
		{URL: ts.URL + "/one"}, {URL: "http://internal.test/"},
	})
	if err != nil || reused {
		t.Fatalf("submit: %v %v", err, reused)
	}
	if j := waitJob(t, s.Batch, id); j.Status != "partial" {
		t.Fatalf("job %+v", j)
	}
	_, ts2, _ := s.Batch.Status(id)
	if ts2[0].Status != "done" || ts2[0].Result == nil || !strings.Contains(ts2[0].Result.Content, "hello batch /one") {
		t.Fatalf("target 0: %+v", ts2[0])
	}
	if ts2[0].Result.Meta["mode"] != "clean" || ts2[0].Result.Meta["untrusted"] != true {
		t.Fatalf("target 0 meta %v: not the fetchTarget envelope", ts2[0].Result.Meta)
	}
	if ts2[1].Status != "failed" || !strings.Contains(ts2[1].Error, "blocked") {
		t.Fatalf("target 1: %+v", ts2[1])
	}
	out := promText(t, s)
	for _, want := range []string{
		`sift_batch_jobs_total{outcome="created"} 1`,
		`sift_batch_jobs_total{outcome="partial"} 1`,
		`sift_batch_targets_total{outcome="ok"} 1`,
		`sift_batch_targets_total{outcome="blocked"} 1`,
		`sift_batch_result_bytes_count 1`,
		`sift_batch_queue_depth 0`,
	} {
		if !strings.Contains(out, want) {
			t.Errorf("metrics lack %q", want)
		}
	}
}

// The adapter turns fetch errors into what the runner and the metrics need.
func TestBatchOutcomeClasses(t *testing.T) {
	cases := []struct {
		err  error
		want string
	}{
		{nil, "ok"},
		{fmt.Errorf("busy: %w", batch.ErrRetryLater), "busy"},
		{&batch.TargetError{Status: 403, Msg: "blocked: private address"}, "blocked"},
		{&batch.TargetError{Status: 429, Msg: "busy"}, "busy"},
		{&batch.TargetError{Status: 504, Msg: "timeout"}, "timeout"},
		{&batch.TargetError{Status: 415, Msg: "unsupported content type"}, "unsupported"},
		{&batch.TargetError{Status: 422, Msg: "document too long for question (6001 chunks); use grep or pages"}, "too_long"},
		{&batch.TargetError{Status: 422, Msg: "unreadable PDF"}, "upstream_4xx"},
		{&batch.TargetError{Status: 404, Msg: "not found"}, "upstream_4xx"},
		{&batch.TargetError{Status: 500, Msg: "internal error"}, "upstream_5xx"},
		{errors.New("something else"), "upstream_5xx"},
	}
	for _, c := range cases {
		if got := batchOutcome(c.err); got != c.want {
			t.Errorf("batchOutcome(%v) = %q, want %q", c.err, got, c.want)
		}
	}
}

// batchFetch maps a busy fetch to ErrRetryLater (the runner retries it) and every other
// failure to a TargetError with the /v1/fetch status.
func TestBatchFetchErrors(t *testing.T) {
	s, _ := newBatchServer(t, "test")
	c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
	_, err := s.batchFetch(context.Background(), c, batch.Target{URL: "http://internal.test/"})
	var te *batch.TargetError
	if !errors.As(err, &te) || te.Status != 403 || !strings.Contains(te.Msg, "blocked") {
		t.Fatalf("private name: %v, want a TargetError 403 blocked", err)
	}
	// Fill the global fetch gate so the next fetch is refused busy.
	for i := 0; i < 8; i++ {
		if err := s.acquire(fmt.Sprintf("filler%d", i), "http://filler.test/"); err != nil {
			t.Fatal(err)
		}
	}
	ts := pageServer(t, nil)
	_, err = s.batchFetch(context.Background(), c, batch.Target{URL: ts.URL})
	if !errors.Is(err, batch.ErrRetryLater) {
		t.Fatalf("busy gate: %v, want an error wrapping batch.ErrRetryLater", err)
	}
}