majordomo/provider/llamaswap/segment_restore_test.go

package llamaswap

import (
	"context"
	"errors"
	"math"
	"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)
	}
	for name, bad := range map[string]float64{
		"NaN":  math.NaN(),
		"+Inf": math.Inf(1),
		"-Inf": math.Inf(-1),
	} {
		if _, err := sg.Segment(context.Background(),
			imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}, Prompt: "dog", Threshold: bad}); !errors.Is(err, llm.ErrUnsupported) {
			t.Errorf("threshold %s: err = %v, want ErrUnsupported", name, 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)
	}
}