404 lines
14 KiB
Go
404 lines
14 KiB
Go
package api
|
||
|
||
// P5a task 06: PDFs through /v1/fetch and /v2/scrape — extraction in the child, the pool, the
|
||
// document cache, the map, pages, grep, and the errors.
|
||
//
|
||
// The child is the test binary itself (TestMain below), as in internal/pdftext's tests.
|
||
// SIFT_PDFTEXT_RUNLOG makes every child start append a line to a file, so a test can count
|
||
// extractions; SIFT_PDFTEXT_SLEEP slows a child down. Both hooks live only in this file.
|
||
|
||
import (
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"log/slog"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"net/netip"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
|
||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||
"git.wntrmute.dev/kyle/sift/internal/digest"
|
||
"git.wntrmute.dev/kyle/sift/internal/fetch"
|
||
"git.wntrmute.dev/kyle/sift/internal/guard"
|
||
"git.wntrmute.dev/kyle/sift/internal/pdfpool"
|
||
"git.wntrmute.dev/kyle/sift/internal/pdftest"
|
||
"git.wntrmute.dev/kyle/sift/internal/pdftext"
|
||
)
|
||
|
||
func TestMain(m *testing.M) {
|
||
if os.Getenv("SIFT_PDFTEXT_CHILD") == "1" {
|
||
if p := os.Getenv("SIFT_PDFTEXT_RUNLOG"); p != "" {
|
||
f, err := os.OpenFile(p, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
|
||
if err == nil {
|
||
_, _ = f.WriteString("run\n")
|
||
_ = f.Close()
|
||
}
|
||
}
|
||
if d, err := time.ParseDuration(os.Getenv("SIFT_PDFTEXT_SLEEP")); err == nil {
|
||
time.Sleep(d)
|
||
}
|
||
os.Exit(pdftext.ChildMain(os.Args[1:], os.Stdin, os.Stdout))
|
||
}
|
||
os.Exit(m.Run())
|
||
}
|
||
|
||
func threePDF() []byte {
|
||
return pdftest.Build(pdftest.Doc{
|
||
Title: "Three Pages",
|
||
Pages: []string{"Chapter one begins here.", "The quick brown fox is on page two.", "Last page."},
|
||
Outline: []pdftest.Outline{{Title: "1 Intro", Page: 1, Level: 1}, {Title: "2 Fox", Page: 2, Level: 1}},
|
||
})
|
||
}
|
||
|
||
func longPDF(n int) []byte {
|
||
var pages []string
|
||
for p := 1; p <= n; p++ {
|
||
pages = append(pages, pdftest.Lorem(p, 40))
|
||
}
|
||
return pdftest.Build(pdftest.Doc{Title: "Long Thesis", Pages: pages, Outline: []pdftest.Outline{{Title: "1 Introduction", Page: 1, Level: 1}, {Title: "2 Method", Page: n / 2, Level: 1}}})
|
||
}
|
||
|
||
// origin serves one PDF and counts GETs. It honours If-None-Match against its current ETag.
|
||
type origin struct {
|
||
mu sync.Mutex
|
||
ct string
|
||
body []byte
|
||
etag string
|
||
gets atomic.Int32
|
||
srv *httptest.Server
|
||
}
|
||
|
||
func newOrigin(t *testing.T, ct string, body []byte) *origin {
|
||
o := &origin{ct: ct, body: body, etag: `"v1"`}
|
||
o.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
o.gets.Add(1)
|
||
o.mu.Lock()
|
||
defer o.mu.Unlock()
|
||
if r.Header.Get("If-None-Match") == o.etag {
|
||
w.WriteHeader(http.StatusNotModified)
|
||
return
|
||
}
|
||
w.Header().Set("ETag", o.etag)
|
||
if o.ct != "" {
|
||
w.Header().Set("Content-Type", o.ct)
|
||
}
|
||
_, _ = w.Write(o.body)
|
||
}))
|
||
t.Cleanup(o.srv.Close)
|
||
return o
|
||
}
|
||
|
||
func (o *origin) replace(body []byte, etag string) {
|
||
o.mu.Lock()
|
||
o.body, o.etag = body, etag
|
||
o.mu.Unlock()
|
||
}
|
||
|
||
type pdfRig struct {
|
||
s *Server
|
||
fleet string // bearer key of caller "test" (hermes-fleet)
|
||
other string // bearer key of caller "other" (hermes-fleet)
|
||
gs string // bearer key of restricted caller "straylight-gs"
|
||
runlog string
|
||
logs *recHandler
|
||
}
|
||
|
||
func newPDFRig(t *testing.T, childEnv ...string) *pdfRig {
|
||
t.Helper()
|
||
dir := t.TempDir()
|
||
keys := filepath.Join(dir, "keys.toml")
|
||
key := func(name, ws string) string {
|
||
out := filepath.Join(dir, name)
|
||
if err := auth.Add(keys, name, ws, out); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
b, err := os.ReadFile(out)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return strings.TrimSpace(string(b))
|
||
}
|
||
r := &pdfRig{fleet: key("test", "hermes-fleet"), other: key("other", "hermes-fleet"), gs: key("straylight-gs", "guardian-star"), runlog: filepath.Join(dir, "runs"), logs: &recHandler{}}
|
||
store, err := auth.Load(keys)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}}
|
||
s := New("test", g, &fetch.Fetcher{Guard: g, UA: "sift-test"}, store, slog.New(r.logs))
|
||
s.PermitRestricted = []string{"straylight-gs"}
|
||
db, err := dcache.Open(filepath.Join(dir, "cache.db"), 0, nil)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
t.Cleanup(func() { _ = db.Close() })
|
||
s.Digest = &digest.Engine{Cache: db} // the cache lives on the digest engine; no model is configured
|
||
s.PDFRunner = &pdftext.Runner{Exe: os.Args[0], Env: append([]string{"SIFT_PDFTEXT_CHILD=1", "SIFT_PDFTEXT_RUNLOG=" + r.runlog}, childEnv...), Deadline: 20 * time.Second}
|
||
r.s = s
|
||
return r
|
||
}
|
||
|
||
func (r *pdfRig) runs() int {
|
||
b, _ := os.ReadFile(r.runlog)
|
||
return strings.Count(string(b), "run\n")
|
||
}
|
||
|
||
func (r *pdfRig) fetch(t *testing.T, key string, body map[string]any) (int, map[string]any, http.Header) {
|
||
t.Helper()
|
||
rec := postJSON(t, r.s, "/v1/fetch", key, body)
|
||
var out map[string]any
|
||
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
||
return rec.Code, out, rec.Header()
|
||
}
|
||
|
||
func pdfInfo(t *testing.T, out map[string]any) map[string]any {
|
||
t.Helper()
|
||
p, ok := out["pdf"].(map[string]any)
|
||
if !ok {
|
||
t.Fatalf("envelope has no pdf object: %v", out)
|
||
}
|
||
return p
|
||
}
|
||
|
||
func TestPDFCleanFetch(t *testing.T) {
|
||
for _, ct := range []string{"application/pdf", "application/octet-stream", ""} {
|
||
r := newPDFRig(t)
|
||
o := newOrigin(t, ct, threePDF())
|
||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
|
||
if code != 200 {
|
||
t.Fatalf("%q: %d %v", ct, code, out)
|
||
}
|
||
content, _ := out["content"].(string)
|
||
if !strings.Contains(content, "--- page 2 ---") || !strings.Contains(content, "quick brown fox") || out["mode"] != "clean" {
|
||
t.Errorf("%q: content %q mode %v", ct, content, out["mode"])
|
||
}
|
||
p := pdfInfo(t, out)
|
||
if p["pages"] != 3.0 || p["title"] != "Three Pages" || p["map"] != false || p["first"] != 1.0 || p["last"] != 3.0 {
|
||
t.Errorf("%q: pdf = %v", ct, p)
|
||
}
|
||
if ol, _ := p["outline"].([]any); len(ol) != 2 {
|
||
t.Errorf("%q: outline = %v", ct, p["outline"])
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestPDFMap(t *testing.T) {
|
||
r := newPDFRig(t)
|
||
r.s.OutCap = 20000
|
||
o := newOrigin(t, "application/pdf", longPDF(40)) // ~40 × 2,700 characters
|
||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
|
||
if code != 200 {
|
||
t.Fatalf("%d %v", code, out)
|
||
}
|
||
content, _ := out["content"].(string)
|
||
p := pdfInfo(t, out)
|
||
if p["map"] != true || !strings.Contains(content, "## Outline") || !strings.Contains(content, "- 2 Method (p. 20)") || len([]rune(content)) > 20000 {
|
||
t.Errorf("map: pdf %v, %d chars:\n%.400s", p, len([]rune(content)), content)
|
||
}
|
||
}
|
||
|
||
func TestPDFPagesAndGrep(t *testing.T) {
|
||
r := newPDFRig(t)
|
||
o := newOrigin(t, "application/pdf", threePDF())
|
||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "pages": "2-3"})
|
||
content, _ := out["content"].(string)
|
||
if code != 200 || strings.Contains(content, "--- page 1 ---") || !strings.Contains(content, "--- page 3 ---") {
|
||
t.Errorf("pages 2-3: %d %q", code, content)
|
||
}
|
||
if p := pdfInfo(t, out); p["first"] != 2.0 || p["last"] != 3.0 {
|
||
t.Errorf("pages 2-3: pdf %v", p)
|
||
}
|
||
for _, bad := range []map[string]any{{"pages": "0-1"}, {"pages": "9"}, {"grep": strings.Repeat("a", 201)}, {"grep": "(unclosed"}} {
|
||
bad["url"] = o.srv.URL
|
||
if code, out, _ := r.fetch(t, r.fleet, bad); code != 400 || out["success"] != false {
|
||
t.Errorf("%v: %d %v, want 400", bad, code, out)
|
||
}
|
||
}
|
||
code, out, _ = r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "grep": "quick brown"})
|
||
content, _ = out["content"].(string)
|
||
if code != 200 || !strings.Contains(content, "p. 2: The quick brown fox is on page two.") || pdfInfo(t, out)["grep_hits"] != 1.0 {
|
||
t.Errorf("grep: %d %q %v", code, content, out["pdf"])
|
||
}
|
||
// pages and grep on a non-PDF are refused rather than ignored.
|
||
html := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||
w.Header().Set("Content-Type", "text/html")
|
||
_, _ = io.WriteString(w, "<html><body><p>hi</p></body></html>")
|
||
}))
|
||
defer html.Close()
|
||
if code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": html.URL, "pages": "1"}); code != 400 {
|
||
t.Errorf("pages on html: %d %v, want 400", code, out)
|
||
}
|
||
}
|
||
|
||
func TestPDFCache(t *testing.T) {
|
||
r := newPDFRig(t)
|
||
o := newOrigin(t, "application/pdf", threePDF())
|
||
get := func(extra map[string]any) string {
|
||
body := map[string]any{"url": o.srv.URL}
|
||
for k, v := range extra {
|
||
body[k] = v
|
||
}
|
||
code, out, _ := r.fetch(t, r.fleet, body)
|
||
if code != 200 {
|
||
t.Fatalf("%v: %d %v", extra, code, out)
|
||
}
|
||
c, _ := out["content"].(string)
|
||
return c
|
||
}
|
||
get(nil)
|
||
if o.gets.Load() != 1 || r.runs() != 1 {
|
||
t.Fatalf("first fetch: %d GETs, %d extractions; want 1 and 1", o.gets.Load(), r.runs())
|
||
}
|
||
if c := get(map[string]any{"pages": "2"}); !strings.Contains(c, "quick brown fox") {
|
||
t.Errorf("cached pages=2: %q", c)
|
||
}
|
||
if o.gets.Load() != 1 || r.runs() != 1 {
|
||
t.Errorf("fresh cache hit: %d GETs, %d extractions; want still 1 and 1", o.gets.Load(), r.runs())
|
||
}
|
||
// Stale: a conditional GET, answered 304, re-uses the text.
|
||
r.s.PDFTTL = time.Nanosecond
|
||
time.Sleep(2 * time.Millisecond)
|
||
get(map[string]any{"grep": "Chapter"})
|
||
if o.gets.Load() != 2 || r.runs() != 1 {
|
||
t.Errorf("after 304: %d GETs, %d extractions; want 2 and 1", o.gets.Load(), r.runs())
|
||
}
|
||
// The document changed: a full GET and a new extraction.
|
||
o.replace(pdftest.Build(pdftest.Doc{Pages: []string{"Brand new text."}}), `"v2"`)
|
||
if c := get(nil); !strings.Contains(c, "Brand new text.") {
|
||
t.Errorf("changed document: %q", c)
|
||
}
|
||
if o.gets.Load() != 3 || r.runs() != 2 {
|
||
t.Errorf("after change: %d GETs, %d extractions; want 3 and 2", o.gets.Load(), r.runs())
|
||
}
|
||
if DefaultPDFTTL != 24*time.Hour {
|
||
t.Errorf("DefaultPDFTTL = %v, want 24h", DefaultPDFTTL)
|
||
}
|
||
}
|
||
|
||
func TestPDFErrors(t *testing.T) {
|
||
r := newPDFRig(t)
|
||
enc, err := os.ReadFile(filepath.Join("..", "..", "tests", "testdata", "pdf", "encrypted-user.pdf"))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
for _, tc := range []struct {
|
||
name string
|
||
body []byte
|
||
want string
|
||
hint string
|
||
}{
|
||
{"scanned", pdftest.Build(pdftest.Doc{Pages: []string{"", ""}, Scanned: []int{1, 2}}), "no text layer", "ocr"},
|
||
{"encrypted", enc, "encrypted", ""},
|
||
{"garbage", []byte("%PDF-1.4\nthis is not really a pdf"), "unreadable", ""},
|
||
} {
|
||
o := newOrigin(t, "application/pdf", tc.body)
|
||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
|
||
msg, _ := out["error"].(string)
|
||
if code != 422 || !strings.Contains(msg, tc.want) || (tc.hint != "" && out["hint"] != tc.hint) {
|
||
t.Errorf("%s: %d %v; want 422 %q hint %q", tc.name, code, out, tc.want, tc.hint)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestPDFTimeout(t *testing.T) {
|
||
r := newPDFRig(t, "SIFT_PDFTEXT_SLEEP=10s")
|
||
r.s.PDFRunner.Deadline = 300 * time.Millisecond
|
||
o := newOrigin(t, "application/pdf", threePDF())
|
||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
|
||
if msg, _ := out["error"].(string); code != 422 || !strings.Contains(msg, "timed out") {
|
||
t.Errorf("slow child: %d %v; want 422 … timed out", code, out)
|
||
}
|
||
}
|
||
|
||
func TestPDFBusy(t *testing.T) {
|
||
r := newPDFRig(t, "SIFT_PDFTEXT_SLEEP=1500ms")
|
||
r.s.PDFPool = pdfpool.New(1, 0, 100*time.Millisecond)
|
||
a := newOrigin(t, "application/pdf", threePDF())
|
||
b := newOrigin(t, "application/pdf", longPDF(2))
|
||
done := make(chan int, 1)
|
||
go func() {
|
||
code, _, _ := r.fetch(t, r.fleet, map[string]any{"url": a.srv.URL})
|
||
done <- code
|
||
}()
|
||
time.Sleep(400 * time.Millisecond) // the first extraction holds the only slot
|
||
code, out, hdr := r.fetch(t, r.other, map[string]any{"url": b.srv.URL})
|
||
if code != 503 || hdr.Get("Retry-After") != "10" || !strings.Contains(fmt.Sprint(out["error"]), "busy") {
|
||
t.Errorf("second caller with the pool full: %d %v Retry-After %q; want 503, busy, 10", code, out, hdr.Get("Retry-After"))
|
||
}
|
||
code, out, _ = r.fetch(t, r.fleet, map[string]any{"url": b.srv.URL})
|
||
if code != 429 {
|
||
t.Errorf("same caller, second document while its first is running: %d %v; want 429", code, out)
|
||
}
|
||
if c := <-done; c != 200 {
|
||
t.Errorf("first fetch: %d", c)
|
||
}
|
||
}
|
||
|
||
func TestPDFRestrictedCaller(t *testing.T) {
|
||
r := newPDFRig(t)
|
||
o := newOrigin(t, "application/pdf", threePDF())
|
||
if code, out, _ := r.fetch(t, r.gs, map[string]any{"url": o.srv.URL}); code != 200 {
|
||
t.Fatalf("restricted fetch: %d %v", code, out)
|
||
}
|
||
r.logs.mu.Lock()
|
||
for _, rec := range r.logs.recs {
|
||
var b strings.Builder
|
||
b.WriteString(rec.Message)
|
||
rec.Attrs(func(a slog.Attr) bool { fmt.Fprintf(&b, " %s=%v", a.Key, a.Value); return true })
|
||
line := b.String()
|
||
if strings.Contains(line, "127.0.0.1") || strings.Contains(line, "Three Pages") || strings.Contains(line, "quick brown") {
|
||
t.Errorf("a restricted caller's PDF fetch logged %q", line)
|
||
}
|
||
}
|
||
r.logs.mu.Unlock()
|
||
// The fleet does not see the restricted caller's cached text: it extracts its own copy.
|
||
if code, _, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL}); code != 200 || r.runs() != 2 {
|
||
t.Errorf("fleet after restricted: %d, %d extractions; want 200 and 2 (separate cache partitions)", code, r.runs())
|
||
}
|
||
}
|
||
|
||
func TestPDFScrape(t *testing.T) {
|
||
r := newPDFRig(t)
|
||
o := newOrigin(t, "application/pdf", threePDF())
|
||
rec := postJSON(t, r.s, "/v2/scrape", r.fleet, map[string]any{"url": o.srv.URL, "formats": []string{"markdown"}})
|
||
var out struct {
|
||
Success bool `json:"success"`
|
||
Data struct {
|
||
Markdown string `json:"markdown"`
|
||
Metadata map[string]any `json:"metadata"`
|
||
} `json:"data"`
|
||
}
|
||
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || rec.Code != 200 || !out.Success {
|
||
t.Fatalf("%d %s", rec.Code, rec.Body.String())
|
||
}
|
||
if !strings.Contains(out.Data.Markdown, "--- page 2 ---") || out.Data.Metadata["numPages"] != 3.0 || out.Data.Metadata["title"] != "Three Pages" {
|
||
t.Errorf("scrape: markdown %q metadata %v", out.Data.Markdown, out.Data.Metadata)
|
||
}
|
||
if ct, _ := out.Data.Metadata["contentType"].(string); !strings.HasPrefix(ct, "application/pdf") {
|
||
t.Errorf("contentType = %q", ct)
|
||
}
|
||
}
|
||
|
||
// A question on a document longer than PDFMaxQuestionChunks is refused with a pointer to grep
|
||
// and pages, before any model call.
|
||
func TestPDFQuestionTooLong(t *testing.T) {
|
||
r := newPDFRig(t)
|
||
r.s.PDFMaxQuestionChunks = 2
|
||
o := newOrigin(t, "application/pdf", longPDF(12))
|
||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is the method?"})
|
||
if msg, _ := out["error"].(string); code != 422 || !strings.Contains(msg, "too long") || !strings.Contains(msg, "grep") {
|
||
t.Errorf("long question: %d %v; want 422 … too long … grep", code, out)
|
||
}
|
||
// P5b: long documents are answered by retrieval now; the limit is only a safety cap.
|
||
if DefaultPDFMaxQuestionChunks != 5000 {
|
||
t.Errorf("DefaultPDFMaxQuestionChunks = %d, want 5000", DefaultPDFMaxQuestionChunks)
|
||
}
|
||
}
|