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

package api

import (
	"encoding/json"
	"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"
)

// newBatchServer builds a Server with the batch runner wired over an in-memory
// dcache, exactly as serve does in production. callers are the caller names to
// provision (each with a bearer key).
func newBatchServer(t *testing.T, g *guard.Guard, 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)
	}
	if g == nil {
		g = &guard.Guard{}
	}
	f := &fetch.Fetcher{Guard: g, UA: "sift-test"}
	s := New("test", g, f, store, slog.New(slog.NewTextHandler(io.Discard, nil)))
	db, err := dcache.Open("", 1<<20, slog.New(slog.NewTextHandler(io.Discard, nil)))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })
	s.Batch = batch.New(db, s.fetchTarget, batch.Config{}, slog.New(slog.NewTextHandler(io.Discard, nil)))
	return s, keys
}

func submitBatch(t *testing.T, s *Server, key string, targets []map[string]string) *httptest.ResponseRecorder {
	t.Helper()
	return postJSON(t, s, "/v1/batch", key, map[string]any{"targets": targets})
}

func getBatch(t *testing.T, s *Server, key, id string) *httptest.ResponseRecorder {
	t.Helper()
	req := httptest.NewRequest(http.MethodGet, "/v1/batch/"+id, nil)
	req.Header.Set("Authorization", "Bearer "+key)
	rec := httptest.NewRecorder()
	s.ServeHTTP(rec, req)
	return rec
}

func batchStatus(t *testing.T, s *Server, key, id string) map[string]any {
	t.Helper()
	rec := getBatch(t, s, key, id)
	var env map[string]any
	if err := json.Unmarshal(rec.Body.Bytes(), &env); err != nil {
		t.Fatalf("status decode: %v (%s)", err, rec.Body.String())
	}
	return env
}

// waitBatchTerminal polls until the job reaches a terminal status.
func waitBatchTerminal(t *testing.T, s *Server, key, id string) map[string]any {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		env := batchStatus(t, s, key, id)
		if st, _ := env["status"].(string); st == "done" || st == "partial" || st == "failed" || st == "cancelled" {
			return env
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatal("batch never reached a terminal status")
	return nil
}

func TestBatchSubmitAndCollect(t *testing.T) {
	g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		_, _ = w.Write([]byte(`<html><head><title>T</title></head><body><article><p>hello batch</p></article></body></html>`))
	}))
	defer ts.Close()
	s, keys := newBatchServer(t, g, "test")
	rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": ts.URL}, {"url": ts.URL + "/2"}})
	if rec.Code != 202 {
		t.Fatalf("submit %d %s", rec.Code, rec.Body.String())
	}
	var sub map[string]any
	if err := json.Unmarshal(rec.Body.Bytes(), &sub); err != nil {
		t.Fatal(err)
	}
	if sub["success"] != true {
		t.Fatalf("success %v", sub["success"])
	}
	id, _ := sub["job_id"].(string)
	if id == "" {
		t.Fatal("missing job_id")
	}
	if reused, _ := sub["reused"].(bool); reused {
		t.Fatal("first submit must not be reused")
	}
	env := waitBatchTerminal(t, s, keys["test"], id)
	if env["status"] != "done" {
		t.Fatalf("status %v", env["status"])
	}
	targets, _ := env["targets"].([]any)
	if len(targets) != 2 {
		t.Fatalf("targets %d", len(targets))
	}
	first := targets[0].(map[string]any)
	if first["status"] != "done" {
		t.Fatalf("target0 %v", first["status"])
	}
	res, _ := first["result"].(map[string]any)
	if res == nil || !strings.Contains(res["content"].(string), "hello batch") {
		t.Fatalf("target0 result %v", res)
	}
}

func TestBatchIdempotentResubmit(t *testing.T) {
	g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		_, _ = w.Write([]byte("ok"))
	}))
	defer ts.Close()
	s, keys := newBatchServer(t, g, "test")
	targets := []map[string]string{{"url": ts.URL}}
	first := submitBatch(t, s, keys["test"], targets)
	second := submitBatch(t, s, keys["test"], targets)
	var a, b map[string]any
	_ = json.Unmarshal(first.Body.Bytes(), &a)
	_ = json.Unmarshal(second.Body.Bytes(), &b)
	if a["job_id"] != b["job_id"] {
		t.Fatalf("ids differ: %v vs %v", a["job_id"], b["job_id"])
	}
	if reused, _ := b["reused"].(bool); !reused {
		t.Fatal("resubmit must be reused=true")
	}
}

func TestBatchGuardRefusedUpFront(t *testing.T) {
	s, keys := newBatchServer(t, nil, "test")
	rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": "http://example.com:8080/"}})
	if rec.Code != 422 {
		t.Fatalf("want 422, got %d %s", rec.Code, rec.Body.String())
	}
	if !strings.Contains(rec.Body.String(), "blocked") {
		t.Fatalf("body %s", rec.Body.String())
	}
}

func TestBatchTooManyTargets(t *testing.T) {
	s, keys := newBatchServer(t, nil, "test")
	var targets []map[string]string
	for i := 0; i < 33; i++ {
		targets = append(targets, map[string]string{"url": "https://example.com/"})
	}
	rec := submitBatch(t, s, keys["test"], targets)
	if rec.Code != 422 {
		t.Fatalf("want 422, got %d %s", rec.Code, rec.Body.String())
	}
}

func TestBatchPartialFailure(t *testing.T) {
	g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		_, _ = w.Write([]byte("ok"))
	}))
	defer ts.Close()
	s, keys := newBatchServer(t, g, "test")
	// target 1 is a private address: passes the submit pre-check (valid URL) but
	// is blocked by the connect-time guard, surfacing as a per-target failure.
	rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": ts.URL}, {"url": "http://10.0.0.5/"}})
	if rec.Code != 202 {
		t.Fatalf("submit %d %s", rec.Code, rec.Body.String())
	}
	var sub map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &sub)
	id := sub["job_id"].(string)
	env := waitBatchTerminal(t, s, keys["test"], id)
	if env["status"] != "partial" {
		t.Fatalf("status %v", env["status"])
	}
	targets := env["targets"].([]any)
	second := targets[1].(map[string]any)
	if second["status"] != "failed" {
		t.Fatalf("target1 status %v", second["status"])
	}
	if err, _ := second["error"].(string); !strings.Contains(err, "blocked") {
		t.Fatalf("target1 error %q", err)
	}
}

func TestBatchOwnership404(t *testing.T) {
	s, keys := newBatchServer(t, nil, "test", "other")
	rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": "https://example.com/"}})
	if rec.Code != 202 {
		t.Fatalf("submit %d", rec.Code)
	}
	var sub map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &sub)
	id := sub["job_id"].(string)
	if got := getBatch(t, s, keys["other"], id); got.Code != 404 {
		t.Fatalf("other caller GET: %d (must not reveal the job)", got.Code)
	}
	if got := getBatch(t, s, keys["test"], id); got.Code != 200 {
		t.Fatalf("owner GET: %d", got.Code)
	}
}

func TestBatchCancel(t *testing.T) {
	g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}}
	s, keys := newBatchServer(t, g, "test")
	rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": "https://example.com/"}})
	var sub map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &sub)
	id := sub["job_id"].(string)
	req := httptest.NewRequest(http.MethodDelete, "/v1/batch/"+id, nil)
	req.Header.Set("Authorization", "Bearer "+keys["test"])
	drec := httptest.NewRecorder()
	s.ServeHTTP(drec, req)
	if drec.Code != 200 {
		t.Fatalf("cancel %d %s", drec.Code, drec.Body.String())
	}
	env := waitBatchTerminal(t, s, keys["test"], id)
	if env["status"] != "cancelled" {
		t.Fatalf("status %v", env["status"])
	}
}