package llamaswap import ( "context" "encoding/base64" "errors" "net/http" "net/http/httptest" "strings" "testing" "gitea.stevedudenhoeffer.com/steve/majordomo/imagegen" "gitea.stevedudenhoeffer.com/steve/majordomo/llm" "gitea.stevedudenhoeffer.com/steve/majordomo/videogen" ) func pngFixture(t *testing.T) []byte { t.Helper() raw, err := base64.StdEncoding.DecodeString(onePixelPNG) if err != nil { t.Fatalf("decode fixture: %v", err) } return raw } func TestUpscale(t *testing.T) { png := pngFixture(t) var gotPath, gotScale, 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) } gotScale = r.FormValue("scale") 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())) up, err := p.UpscaleModel("mediautils") if err != nil { t.Fatalf("UpscaleModel: %v", err) } res, err := up.Upscale(context.Background(), imagegen.UpscaleRequest{Image: imagegen.Image{MIME: "image/png", Data: png}}, imagegen.WithUpscaleScale(2)) if err != nil { t.Fatalf("Upscale: %v", err) } if gotPath != "/upstream/mediautils/v1/upscale" { t.Errorf("path = %q", gotPath) } if gotScale != "2" || gotFilename != "image.png" { t.Errorf("scale/filename = %q/%q", gotScale, gotFilename) } if len(res.Images) != 1 || res.Images[0].MIME != "image/png" { t.Fatalf("images = %+v", res.Images) } } func TestUpscaleRejectsBadArgs(t *testing.T) { p := New(WithBaseURL("http://unused")) up, _ := p.UpscaleModel("mediautils") if _, err := up.Upscale(context.Background(), imagegen.UpscaleRequest{}); !errors.Is(err, llm.ErrUnsupported) { t.Errorf("no image: err = %v, want ErrUnsupported", err) } if _, err := up.Upscale(context.Background(), imagegen.UpscaleRequest{Image: imagegen.Image{Data: []byte{1}}, Scale: 3}); !errors.Is(err, llm.ErrUnsupported) { t.Errorf("scale 3: err = %v, want ErrUnsupported", err) } } func TestUpscaleRejectsNonImageResponse(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())) up, _ := p.UpscaleModel("mediautils") _, err := up.Upscale(context.Background(), imagegen.UpscaleRequest{Image: imagegen.Image{Data: pngFixture(t)}}) var apiErr *llm.APIError if !errors.As(err, &apiErr) { t.Fatalf("err = %v, want APIError for non-image body", err) } } func TestRemoveBackground(t *testing.T) { png := pngFixture(t) var gotPath, gotNet, gotOM 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) } gotNet = r.FormValue("model") gotOM = r.FormValue("om") w.Header().Set("Content-Type", "image/png") _, _ = w.Write(png) })) defer srv.Close() p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) br, err := p.BackgroundRemovalModel("rembg") if err != nil { t.Fatalf("BackgroundRemovalModel: %v", err) } res, err := br.RemoveBackground(context.Background(), imagegen.BackgroundRemovalRequest{Image: imagegen.Image{Data: png}}, imagegen.WithBackgroundNet("birefnet-general"), imagegen.WithOnlyMask()) if err != nil { t.Fatalf("RemoveBackground: %v", err) } if gotPath != "/upstream/rembg/api/remove" { t.Errorf("path = %q", gotPath) } if gotNet != "birefnet-general" || gotOM != "true" { t.Errorf("net/om = %q/%q", gotNet, gotOM) } if len(res.Images) != 1 { t.Fatalf("images = %+v", res.Images) } } func TestRemoveBackgroundOmitsUnsetFields(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["model"]; ok { t.Error("model field sent for default net; want omitted") } if _, ok := r.MultipartForm.Value["om"]; ok { t.Error("om field sent for cutout mode; want omitted") } w.Header().Set("Content-Type", "image/png") _, _ = w.Write(png) })) defer srv.Close() p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) br, _ := p.BackgroundRemovalModel("rembg") if _, err := br.RemoveBackground(context.Background(), imagegen.BackgroundRemovalRequest{Image: imagegen.Image{Data: png}}); err != nil { t.Fatalf("RemoveBackground: %v", err) } } // mp4Fixture is a minimal ftyp box so http.DetectContentType sniffs video/mp4. func mp4Fixture() []byte { return []byte{0, 0, 0, 24, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm', 0, 0, 2, 0, 'i', 's', 'o', 'm', 'i', 's', 'o', '2'} } func TestInterpolate(t *testing.T) { var gotPath, gotMulti, gotMode, gotTargetFPS 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) } gotMulti = r.FormValue("multi") gotMode = r.FormValue("mode") gotTargetFPS = r.FormValue("target_fps") w.Header().Set("Content-Type", "video/mp4") _, _ = w.Write(mp4Fixture()) })) defer srv.Close() p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) ip, err := p.InterpolatorModel("mediautils") if err != nil { t.Fatalf("InterpolatorModel: %v", err) } res, err := ip.Interpolate(context.Background(), videogen.InterpolateRequest{Video: videogen.Video{Data: mp4Fixture(), MIME: "video/mp4"}}, videogen.WithInterpolateMulti(2), videogen.WithTargetFPS(60)) if err != nil { t.Fatalf("Interpolate: %v", err) } if gotPath != "/upstream/mediautils/v1/interpolate" { t.Errorf("path = %q", gotPath) } if gotMulti != "2" || gotMode != "" || gotTargetFPS != "60" { t.Errorf("multi/mode/target_fps = %q/%q/%q", gotMulti, gotMode, gotTargetFPS) } if res.Video.MIME != "video/mp4" || len(res.Video.Data) == 0 { t.Fatalf("video = %q/%d bytes", res.Video.MIME, len(res.Video.Data)) } } func TestInterpolateSlowMoSetsModeAndDropsTargetFPS(t *testing.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 got := r.FormValue("mode"); got != "slowmo" { t.Errorf("mode = %q, want slowmo", got) } if _, ok := r.MultipartForm.Value["target_fps"]; ok { t.Error("target_fps sent in slowmo mode; want omitted") } w.Header().Set("Content-Type", "video/mp4") _, _ = w.Write(mp4Fixture()) })) defer srv.Close() p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) ip, _ := p.InterpolatorModel("mediautils") _, err := ip.Interpolate(context.Background(), videogen.InterpolateRequest{Video: videogen.Video{Data: mp4Fixture()}, TargetFPS: 60}, videogen.WithSlowMo()) if err != nil { t.Fatalf("Interpolate: %v", err) } } func TestInterpolateRejectsBadArgs(t *testing.T) { p := New(WithBaseURL("http://unused")) ip, _ := p.InterpolatorModel("mediautils") if _, err := ip.Interpolate(context.Background(), videogen.InterpolateRequest{}); !errors.Is(err, llm.ErrUnsupported) { t.Errorf("no video: err = %v, want ErrUnsupported", err) } if _, err := ip.Interpolate(context.Background(), videogen.InterpolateRequest{Video: videogen.Video{Data: []byte{1}}, Multi: 3}); !errors.Is(err, llm.ErrUnsupported) { t.Errorf("multi 3: err = %v, want ErrUnsupported", err) } } func TestUpstreamPathRejectsSeparators(t *testing.T) { for _, bad := range []string{"", "a/b", "a?b", "a#b"} { if _, err := upstreamPath(bad, "/x"); err == nil { t.Errorf("upstreamPath(%q) succeeded; want error", bad) } } got, err := upstreamPath("rembg", "api/remove") if err != nil || got != "/upstream/rembg/api/remove" { t.Errorf("upstreamPath = %q, %v", got, err) } if !strings.HasPrefix(got, "/upstream/") { t.Errorf("path prefix wrong: %q", got) } } func TestUpscaleRejectsEmptyContentTypeNonImage(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header()["Content-Type"] = nil // suppress auto-detection: NO Content-Type at all _, _ = w.Write([]byte("502 bad gateway")) })) defer srv.Close() p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) up, _ := p.UpscaleModel("mediautils") _, err := up.Upscale(context.Background(), imagegen.UpscaleRequest{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) } } func TestUpstreamPathRejectsDotDot(t *testing.T) { if _, err := upstreamPath("..", "/x"); err == nil { t.Error("model '..' accepted") } if _, err := upstreamPath("m", "/v1/audio?path=../../api/models/unload"); err == nil { t.Error("dot-dot rest accepted") } if _, err := upstreamPath("m", "/v1/audio?path=https://evil.example/x"); err == nil { t.Error("absolute-URL rest accepted") } }