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