sift/internal/api/batch_test.go

package api

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

// Every test fetches only from local httptest servers or internal.test (fake DNS): no test
// touches the network. newBatchServer, pageServer and fakeDNS come from batchrun_test.go.

func batchReq(t *testing.T, s *Server, method, path, key string, body any) *httptest.ResponseRecorder {
	t.Helper()
	var rd io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			t.Fatal(err)
		}
		rd = strings.NewReader(string(b))
	}
	req := httptest.NewRequest(method, path, rd)
	req.Header.Set("Authorization", "Bearer "+key)
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	rec := httptest.NewRecorder()
	s.ServeHTTP(rec, req)
	return rec
}

func batchDecode(t *testing.T, rec *httptest.ResponseRecorder) map[string]any {
	t.Helper()
	var m map[string]any
	if err := json.Unmarshal(rec.Body.Bytes(), &m); err != nil {
		t.Fatalf("decode: %v (%d %s)", err, rec.Code, rec.Body.String())
	}
	return m
}

func submit(t *testing.T, s *Server, key string, targets ...map[string]any) map[string]any {
	t.Helper()
	rec := batchReq(t, s, http.MethodPost, "/v1/batch", key, map[string]any{"targets": targets})
	if rec.Code != 202 {
		t.Fatalf("submit: %d %s", rec.Code, rec.Body.String())
	}
	return batchDecode(t, rec)
}

func waitTerminal(t *testing.T, s *Server, key, id string) map[string]any {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		env := batchDecode(t, batchReq(t, s, http.MethodGet, "/v1/batch/"+id, key, nil))
		switch env["status"] {
		case "done", "partial", "failed", "cancelled":
			return env
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatal("batch never reached a terminal status")
	return nil
}

func TestBatchSubmitAndCollect(t *testing.T) {
	ts := pageServer(t, nil)
	s, keys := newBatchServer(t, "test")
	sub := submit(t, s, keys["test"], map[string]any{"url": ts.URL + "/one"}, map[string]any{"url": ts.URL + "/two"})
	id, _ := sub["job_id"].(string)
	if sub["success"] != true || id == "" || sub["reused"] != false || sub["targets"] != float64(2) || sub["status"] != "queued" {
		t.Fatalf("submit body %v", sub)
	}

	// The status body is a summary: statuses, no results.
	env := waitTerminal(t, s, keys["test"], id)
	if env["status"] != "done" || env["targets_total"] != float64(2) || env["targets_done"] != float64(2) {
		t.Fatalf("status %v", env)
	}
	if _, err := time.Parse(time.RFC3339, fmt.Sprint(env["created_at"])); err != nil {
		t.Fatalf("created_at %v is not RFC 3339", env["created_at"])
	}
	targets := env["targets"].([]any)
	if len(targets) != 2 {
		t.Fatalf("targets %d", len(targets))
	}
	for i, x := range targets {
		tg := x.(map[string]any)
		if tg["idx"] != float64(i) || tg["status"] != "done" {
			t.Fatalf("target %d: %v", i, tg)
		}
		if _, has := tg["result"]; has {
			t.Fatalf("summary must not carry results: target %d has one", i)
		}
	}

	// One target's result: the /v1/fetch envelope.
	rec := batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/1", keys["test"], nil)
	if rec.Code != 200 {
		t.Fatalf("GET idx: %d %s", rec.Code, rec.Body.String())
	}
	one := batchDecode(t, rec)
	res, _ := one["result"].(map[string]any)
	if one["idx"] != float64(1) || one["status"] != "done" || res == nil {
		t.Fatalf("GET idx body %v", one)
	}
	if c, _ := res["content"].(string); !strings.Contains(c, "hello batch /two") {
		t.Fatalf("result content %q", c)
	}
	if res["untrusted"] != true || res["mode"] != "clean" {
		t.Fatalf("result is not the /v1/fetch envelope: %v", res)
	}
	if src, _ := res["source"].(map[string]any); src == nil || src["final_url"] != ts.URL+"/two" {
		t.Fatalf("result source %v", res["source"])
	}

	// ?results=1 inlines every result.
	full := batchDecode(t, batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"?results=1", keys["test"], nil))
	for i, x := range full["targets"].([]any) {
		if _, has := x.(map[string]any)["result"]; !has {
			t.Fatalf("results=1: target %d has no result", i)
		}
	}

	if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/2", keys["test"], nil).Code; got != 404 {
		t.Fatalf("GET idx out of range: %d, want 404", got)
	}
	if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/x", keys["test"], nil).Code; got != 404 {
		t.Fatalf("GET idx not a number: %d, want 404", got)
	}
}

func TestBatchIdempotentResubmit(t *testing.T) {
	ts := pageServer(t, nil)
	s, keys := newBatchServer(t, "test")
	tg := map[string]any{"url": ts.URL, "mode": "clean"}
	a := submit(t, s, keys["test"], tg)
	b := submit(t, s, keys["test"], tg)
	if a["job_id"] != b["job_id"] || b["reused"] != true {
		t.Fatalf("resubmit: %v then %v", a, b)
	}
}

func TestBatchRefusedUpFront(t *testing.T) {
	s, keys := newBatchServer(t, "test")
	cases := []struct {
		name string
		body any
		code int
		want string
	}{
		{"guard pre-check (port)", map[string]any{"targets": []any{map[string]any{"url": "http://127.0.0.1/"}, map[string]any{"url": "http://example.com:8080/"}}}, 422, "blocked"},
		{"guard pre-check (private literal)", map[string]any{"targets": []any{map[string]any{"url": "http://10.0.0.5/"}}}, 422, "blocked"},
		{"no targets", map[string]any{"targets": []any{}}, 400, ""},
		{"not json", "nope", 400, ""},
		{"empty url", map[string]any{"targets": []any{map[string]any{"url": ""}}}, 400, ""},
		{"bad mode", map[string]any{"targets": []any{map[string]any{"url": "http://127.0.0.1/", "mode": "shout"}}}, 400, ""},
	}
	for _, c := range cases {
		rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], c.body)
		if rec.Code != c.code || !strings.Contains(rec.Body.String(), c.want) {
			t.Errorf("%s: %d %s, want %d containing %q", c.name, rec.Code, rec.Body.String(), c.code, c.want)
		}
	}
	// The guard refusal names the offending index.
	rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], cases[0].body)
	if !strings.Contains(rec.Body.String(), "1") {
		t.Errorf("guard refusal does not name target 1: %s", rec.Body.String())
	}

	var many []any
	for i := 0; i < 33; i++ {
		many = append(many, map[string]any{"url": fmt.Sprintf("http://127.0.0.1/%d", i)})
	}
	if rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], map[string]any{"targets": many}); rec.Code != 422 {
		t.Errorf("33 targets: %d, want 422", rec.Code)
	}
}

// 32 targets, each with a 500-rune question, fit the batch body cap even though they are well
// over the 64 KiB cap of the other endpoints.
func TestBatchBodyCap(t *testing.T) {
	hold := make(chan struct{})
	defer close(hold)
	ts := pageServer(t, hold)
	s, keys := newBatchServer(t, "test")
	q := strings.Repeat("q", 500)
	var many []any
	for i := 0; i < 32; i++ {
		many = append(many, map[string]any{"url": fmt.Sprintf("%s/%d?pad=%s", ts.URL, i, strings.Repeat("p", 1500)), "question": q})
	}
	if rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], map[string]any{"targets": many}); rec.Code != 202 {
		t.Fatalf("32 targets with questions (~%d KiB): %d %s", 32*2100/1024, rec.Code, rec.Body.String())
	}
}

func TestBatchPartialFailure(t *testing.T) {
	ts := pageServer(t, nil)
	s, keys := newBatchServer(t, "test")
	// Target 1 is a name that resolves to a private address. Names are not resolved at submit,
	// so it passes the pre-check; the fetch-time guard blocks it: a per-target failure.
	sub := submit(t, s, keys["test"], map[string]any{"url": ts.URL}, map[string]any{"url": "http://internal.test/"})
	env := waitTerminal(t, s, keys["test"], sub["job_id"].(string))
	if env["status"] != "partial" {
		t.Fatalf("status %v", env["status"])
	}
	second := env["targets"].([]any)[1].(map[string]any)
	if second["status"] != "failed" {
		t.Fatalf("target 1: %v", second)
	}
	if e, _ := second["error"].(string); !strings.Contains(e, "blocked") {
		t.Fatalf("target 1 error %q", e)
	}
}

func TestBatchOwnership404(t *testing.T) {
	hold := make(chan struct{})
	defer close(hold)
	ts := pageServer(t, hold)
	s, keys := newBatchServer(t, "test", "other")
	id := submit(t, s, keys["test"], map[string]any{"url": ts.URL})["job_id"].(string)
	for _, p := range []string{"/v1/batch/" + id, "/v1/batch/" + id + "/0"} {
		if got := batchReq(t, s, http.MethodGet, p, keys["other"], nil).Code; got != 404 {
			t.Fatalf("other caller GET %s: %d (must not reveal the job)", p, got)
		}
	}
	if got := batchReq(t, s, http.MethodDelete, "/v1/batch/"+id, keys["other"], nil).Code; got != 404 {
		t.Fatalf("other caller DELETE: %d", got)
	}
	if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id, keys["test"], nil).Code; got != 200 {
		t.Fatalf("owner GET: %d", got)
	}
	if got := batchReq(t, s, http.MethodGet, "/v1/batch/no-such-job", keys["test"], nil).Code; got != 404 {
		t.Fatalf("unknown job: %d", got)
	}
	if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id, "not-a-key", nil).Code; got != 401 {
		t.Fatalf("bad key: %d", got)
	}
}

func TestBatchCancel(t *testing.T) {
	hold := make(chan struct{})
	ts := pageServer(t, hold)
	s, keys := newBatchServer(t, "test")
	id := submit(t, s, keys["test"], map[string]any{"url": ts.URL + "/a"}, map[string]any{"url": ts.URL + "/b"})["job_id"].(string)
	rec := batchReq(t, s, http.MethodDelete, "/v1/batch/"+id, keys["test"], nil)
	if rec.Code != 200 {
		t.Fatalf("cancel: %d %s", rec.Code, rec.Body.String())
	}
	if b := batchDecode(t, rec); b["success"] != true || b["status"] != "cancelled" {
		t.Fatalf("cancel body %v", b)
	}
	close(hold)
	env := waitTerminal(t, s, keys["test"], id)
	if env["status"] != "cancelled" {
		t.Fatalf("status %v", env["status"])
	}
	if tg := env["targets"].([]any)[1].(map[string]any); tg["status"] != "pending" {
		t.Fatalf("target 1 after cancel: %v", tg)
	}
	if got := batchReq(t, s, http.MethodDelete, "/v1/batch/"+id, keys["test"], nil).Code; got != 409 {
		t.Fatalf("cancelling a finished job: %d, want 409", got)
	}
}

// The request metrics never get a job id as a label value.
func TestBatchMetricsPathLabel(t *testing.T) {
	ts := pageServer(t, nil)
	s, keys := newBatchServer(t, "test")
	id := submit(t, s, keys["test"], map[string]any{"url": ts.URL})["job_id"].(string)
	waitTerminal(t, s, keys["test"], id)
	batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/0", keys["test"], nil)
	if strings.Contains(promText(t, s), id) {
		t.Fatal("a job id appears in /metrics")
	}
}