majordomo/provider/llamaswap/ocr_test.go

package llamaswap

import (
	"context"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
	"gitea.stevedudenhoeffer.com/steve/majordomo/ocr"
)

func TestRecognize(t *testing.T) {
	var gotPath, gotLangs, gotMaxPages, gotFilename string
	var gotFile []byte
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		if err := r.ParseMultipartForm(1 << 20); err != nil {
			t.Fatalf("parse multipart: %v", err)
		}
		gotLangs = r.FormValue("langs")
		gotMaxPages = r.FormValue("max_pages")
		f, hdr, err := r.FormFile("file")
		if err != nil {
			t.Fatalf("form file: %v", err)
		}
		defer f.Close()
		gotFile, _ = io.ReadAll(f)
		gotFilename = hdr.Filename
		_, _ = w.Write([]byte(`{"pages":[
			{"number":1,"text":"page one","lines":[{"text":"page","bbox":[0,0,1,1],"confidence":0.9},{"text":"one"}],"layout":{}},
			{"number":2,"text":"page two"}
		]}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	om, err := p.OCRModel("ocr-surya")
	if err != nil {
		t.Fatalf("OCRModel: %v", err)
	}
	res, err := om.Recognize(context.Background(),
		ocr.Request{Document: []byte("%PDF"), MIME: "application/pdf"},
		ocr.WithLanguages("en", "de"), ocr.WithMaxPages(5))
	if err != nil {
		t.Fatalf("Recognize: %v", err)
	}
	if gotPath != "/upstream/ocr-surya/v1/ocr" {
		t.Errorf("path = %q", gotPath)
	}
	if gotLangs != "en,de" || gotMaxPages != "5" {
		t.Errorf("langs/max_pages = %q/%q", gotLangs, gotMaxPages)
	}
	if string(gotFile) != "%PDF" || gotFilename != "document.pdf" {
		t.Errorf("file = %q name = %q", gotFile, gotFilename)
	}
	if len(res.Pages) != 2 || res.Pages[0].Number != 1 || res.Pages[0].Text != "page one" ||
		res.Pages[1].Number != 2 || res.Pages[1].Text != "page two" {
		t.Errorf("pages = %+v", res.Pages)
	}
	if res.Text != "page one\n\npage two" {
		t.Errorf("text = %q", res.Text)
	}
	if res.Raw == nil {
		t.Error("Raw = nil, want raw payload")
	}
}

func TestRecognizeJoinsLinesWhenPageTextAbsent(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"pages":[{"lines":[{"text":"first line"},{"text":"second line"}]}]}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	om, _ := p.OCRModel("ocr-surya")
	res, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")})
	if err != nil {
		t.Fatalf("Recognize: %v", err)
	}
	if len(res.Pages) != 1 || res.Pages[0].Text != "first line\nsecond line" {
		t.Errorf("pages = %+v", res.Pages)
	}
	// Missing page number defaults to position.
	if res.Pages[0].Number != 1 {
		t.Errorf("number = %d, want 1", res.Pages[0].Number)
	}
	if res.Text != "first line\nsecond line" {
		t.Errorf("text = %q", res.Text)
	}
}

func TestRecognizeOmitsUnsetFields(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := r.ParseMultipartForm(1 << 20); err != nil {
			t.Fatalf("parse multipart: %v", err)
		}
		for _, k := range []string{"langs", "max_pages"} {
			if v, ok := r.MultipartForm.Value[k]; ok {
				t.Errorf("unset request sent %q = %v, want omitted", k, v)
			}
		}
		if _, hdr, err := r.FormFile("file"); err == nil {
			if hdr.Filename != "document" {
				t.Errorf("filename = %q, want document fallback", hdr.Filename)
			}
		}
		_, _ = w.Write([]byte(`{"pages":[{"number":1,"text":"x"}]}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	om, _ := p.OCRModel("ocr-surya")
	if _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}); err != nil {
		t.Fatalf("Recognize: %v", err)
	}
}

func TestRecognizeRejectsBadArgs(t *testing.T) {
	p := New(WithBaseURL("http://unused"))
	om, _ := p.OCRModel("ocr-surya")
	if _, err := om.Recognize(context.Background(), ocr.Request{}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("no document: err = %v, want ErrUnsupported", err)
	}
	if _, err := om.Recognize(context.Background(),
		ocr.Request{Document: []byte{1}, MaxPages: -1}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("negative max pages: err = %v, want ErrUnsupported", err)
	}
}

func TestRecognizeRejectsZeroPages(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"pages":[]}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	om, _ := p.OCRModel("ocr-surya")
	_, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")})
	var apiErr *llm.APIError
	if !errors.As(err, &apiErr) {
		t.Fatalf("err = %v, want APIError for zero pages", err)
	}
}

func TestRecognizeRejectsNonJSONResponse(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/html")
		_, _ = w.Write([]byte("<html>proxy error page</html>"))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	om, _ := p.OCRModel("ocr-surya")
	if _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}); err == nil ||
		!strings.Contains(err.Error(), "decode ocr response") {
		t.Fatalf("err = %v, want decode error for non-JSON body", err)
	}
}

func TestRecognizeSurfacesAPIError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusBadRequest)
		_, _ = w.Write([]byte(`{"error":{"message":"unsupported file type"}}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	om, _ := p.OCRModel("ocr-surya")
	_, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("bad")})
	var apiErr *llm.APIError
	if !errors.As(err, &apiErr) {
		t.Fatalf("err = %T %v, want *llm.APIError", err, err)
	}
	if apiErr.Status != http.StatusBadRequest || apiErr.Message != "unsupported file type" || apiErr.Model != "ocr-surya" {
		t.Errorf("apiErr = %+v", apiErr)
	}
}

func TestDocumentFilename(t *testing.T) {
	cases := []struct {
		filename, mime, want string
	}{
		{"scan.pdf", "", "scan.pdf"},
		{"evil\r\nname.pdf", "", "evilname.pdf"},
		{"", "application/pdf", "document.pdf"},
		{"", "image/png", "document.png"},
		{"", "image/jpeg", "document.jpg"},
		{"", "image/webp; charset=binary", "document.webp"},
		{"", "", "document"},
	}
	for _, tc := range cases {
		if got := documentFilename(tc.filename, tc.mime); got != tc.want {
			t.Errorf("documentFilename(%q, %q) = %q, want %q", tc.filename, tc.mime, got, tc.want)
		}
	}
}