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("proxy error page")) })) 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) } } }