Files
majordomo/provider/llamaswap/segment_restore_test.go
T
steveandClaude Fable 5 e6987f54b2
CI / Tidy (pull_request) Successful in 9m25s
CI / Build & Test (pull_request) Successful in 9m48s
Gadfly review (reusable) / review (pull_request) Successful in 47m4s
Adversarial Review (Gadfly) / review (pull_request) Successful in 47m4s
feat: wave-3 image + document surfaces — segmentation, colorize, face restore, ocr (ADR-0023)
- imagegen.Segmenter/SegmentationProvider: prompted mask via
  POST /upstream/<id>/v1/segment (file, prompt[, threshold], output=mask);
  white = prompted region, EditRequest.Mask polarity.
- imagegen.Colorizer/ColorizeProvider: POST /upstream/<id>/v1/colorize.
- imagegen.FaceRestorer/FaceRestoreProvider:
  POST /upstream/<id>/v1/restore_faces (upscale 1|2).
- New ocr leaf package (Request/Page/Result, Recognize) + llamaswap
  OCRModel: POST /upstream/<id>/v1/ocr (file[, langs, max_pages]),
  tolerant per-page decode (join lines when page text absent), Raw
  escape hatch.
- httptest contract tests per surface; ADR-0023; ADR index backfilled
  (0020-0022 rows were missing).

Co-Authored-By: Claude Fable 5 <[email protected]>
2026-07-16 16:54:30 -04:00

271 lines
9.1 KiB
Go

package llamaswap
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)
func TestSegment(t *testing.T) {
png := pngFixture(t)
var gotPath, gotPrompt, gotThreshold, gotOutput, gotFilename string
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.Errorf("parse multipart: %v", err)
}
gotPrompt = r.FormValue("prompt")
gotThreshold = r.FormValue("threshold")
gotOutput = r.FormValue("output")
if f, hdr, err := r.FormFile("file"); err == nil {
gotFilename = hdr.Filename
f.Close()
} else {
t.Errorf("file part: %v", err)
}
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
sg, err := p.SegmentationModel("segment-langsam")
if err != nil {
t.Fatalf("SegmentationModel: %v", err)
}
res, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{MIME: "image/png", Data: png}, Prompt: "the red car"},
imagegen.WithSegmentationThreshold(0.35))
if err != nil {
t.Fatalf("Segment: %v", err)
}
if gotPath != "/upstream/segment-langsam/v1/segment" {
t.Errorf("path = %q", gotPath)
}
if gotPrompt != "the red car" || gotThreshold != "0.35" || gotOutput != "mask" || gotFilename != "image.png" {
t.Errorf("prompt/threshold/output/filename = %q/%q/%q/%q", gotPrompt, gotThreshold, gotOutput, gotFilename)
}
if len(res.Images) != 1 || res.Images[0].MIME != "image/png" {
t.Fatalf("images = %+v", res.Images)
}
}
func TestSegmentOmitsDefaultThreshold(t *testing.T) {
png := pngFixture(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Errorf("parse multipart: %v", err)
}
if _, ok := r.MultipartForm.Value["threshold"]; ok {
t.Error("threshold field sent for default; want omitted")
}
if got := r.FormValue("output"); got != "mask" {
t.Errorf("output = %q, want mask (always sent)", got)
}
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
sg, _ := p.SegmentationModel("segment-langsam")
if _, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{Data: png}, Prompt: "dog"}); err != nil {
t.Fatalf("Segment: %v", err)
}
}
func TestSegmentRejectsBadArgs(t *testing.T) {
p := New(WithBaseURL("http://unused"))
sg, _ := p.SegmentationModel("segment-langsam")
if _, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Prompt: "dog"}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("no image: err = %v, want ErrUnsupported", err)
}
if _, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("no prompt: err = %v, want ErrUnsupported", err)
}
if _, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}, Prompt: "dog", Threshold: 1.5}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("threshold 1.5: err = %v, want ErrUnsupported", err)
}
}
func TestSegmentRejectsNonImageResponse(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()))
sg, _ := p.SegmentationModel("segment-langsam")
_, err := sg.Segment(context.Background(),
imagegen.SegmentationRequest{Image: imagegen.Image{Data: pngFixture(t)}, Prompt: "dog"})
var apiErr *llm.APIError
if !errors.As(err, &apiErr) {
t.Fatalf("err = %v, want APIError for non-image body", err)
}
}
func TestColorize(t *testing.T) {
png := pngFixture(t)
var gotPath, gotFilename string
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.Errorf("parse multipart: %v", err)
}
if f, hdr, err := r.FormFile("file"); err == nil {
gotFilename = hdr.Filename
f.Close()
} else {
t.Errorf("file part: %v", err)
}
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
cl, err := p.ColorizeModel("mediautils")
if err != nil {
t.Fatalf("ColorizeModel: %v", err)
}
res, err := cl.Colorize(context.Background(),
imagegen.ColorizeRequest{Image: imagegen.Image{Data: png}})
if err != nil {
t.Fatalf("Colorize: %v", err)
}
if gotPath != "/upstream/mediautils/v1/colorize" {
t.Errorf("path = %q", gotPath)
}
if gotFilename != "image.png" {
t.Errorf("filename = %q", gotFilename)
}
if len(res.Images) != 1 || res.Images[0].MIME != "image/png" {
t.Fatalf("images = %+v", res.Images)
}
}
func TestColorizeRejectsEmptyImage(t *testing.T) {
p := New(WithBaseURL("http://unused"))
cl, _ := p.ColorizeModel("mediautils")
if _, err := cl.Colorize(context.Background(), imagegen.ColorizeRequest{}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("err = %v, want ErrUnsupported", err)
}
}
func TestColorizeSurfacesAPIError(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusServiceUnavailable)
_, _ = w.Write([]byte(`{"error":{"message":"model is loading"}}`))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
cl, _ := p.ColorizeModel("mediautils")
_, err := cl.Colorize(context.Background(),
imagegen.ColorizeRequest{Image: imagegen.Image{Data: pngFixture(t)}})
var apiErr *llm.APIError
if !errors.As(err, &apiErr) {
t.Fatalf("err = %T %v, want *llm.APIError", err, err)
}
if apiErr.Status != http.StatusServiceUnavailable || apiErr.Message != "model is loading" || apiErr.Model != "mediautils" {
t.Errorf("apiErr = %+v", apiErr)
}
}
func TestRestoreFaces(t *testing.T) {
png := pngFixture(t)
var gotPath, gotUpscale string
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.Errorf("parse multipart: %v", err)
}
gotUpscale = r.FormValue("upscale")
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
fr, err := p.FaceRestoreModel("mediautils")
if err != nil {
t.Fatalf("FaceRestoreModel: %v", err)
}
res, err := fr.RestoreFaces(context.Background(),
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: png}},
imagegen.WithFaceRestoreUpscale(2))
if err != nil {
t.Fatalf("RestoreFaces: %v", err)
}
if gotPath != "/upstream/mediautils/v1/restore_faces" {
t.Errorf("path = %q", gotPath)
}
if gotUpscale != "2" {
t.Errorf("upscale = %q", gotUpscale)
}
if len(res.Images) != 1 {
t.Fatalf("images = %+v", res.Images)
}
}
func TestRestoreFacesOmitsDefaultUpscale(t *testing.T) {
png := pngFixture(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Errorf("parse multipart: %v", err)
}
if _, ok := r.MultipartForm.Value["upscale"]; ok {
t.Error("upscale field sent for default; want omitted")
}
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(png)
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
fr, _ := p.FaceRestoreModel("mediautils")
if _, err := fr.RestoreFaces(context.Background(),
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: png}}); err != nil {
t.Fatalf("RestoreFaces: %v", err)
}
}
func TestRestoreFacesRejectsBadArgs(t *testing.T) {
p := New(WithBaseURL("http://unused"))
fr, _ := p.FaceRestoreModel("mediautils")
if _, err := fr.RestoreFaces(context.Background(), imagegen.FaceRestoreRequest{}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("no image: err = %v, want ErrUnsupported", err)
}
if _, err := fr.RestoreFaces(context.Background(),
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: []byte{1}}, Upscale: 3}); !errors.Is(err, llm.ErrUnsupported) {
t.Errorf("upscale 3: err = %v, want ErrUnsupported", err)
}
}
func TestRestoreFacesRejectsNonImageResponse(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header()["Content-Type"] = nil // NO Content-Type at all
_, _ = w.Write([]byte("502 bad gateway"))
}))
defer srv.Close()
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
fr, _ := p.FaceRestoreModel("mediautils")
_, err := fr.RestoreFaces(context.Background(),
imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: pngFixture(t)}})
var apiErr *llm.APIError
if !errors.As(err, &apiErr) {
t.Fatalf("err = %v, want APIError for headerless non-image body", err)
}
}