- 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]>
271 lines
9.1 KiB
Go
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)
|
|
}
|
|
}
|