majordomo/provider/llamaswap/stems_sfx_enhance_test.go

package llamaswap

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"gitea.stevedudenhoeffer.com/steve/majordomo/audio"
	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
	"gitea.stevedudenhoeffer.com/steve/majordomo/musicgen"
)

// stemsZipFixture builds a Demucs-style stems zip: entries under a per-model
// directory, entry name = stem name, extension = container.
func stemsZipFixture(t *testing.T, entries map[string]string) []byte {
	t.Helper()
	var buf bytes.Buffer
	zw := zip.NewWriter(&buf)
	for name, data := range entries {
		w, err := zw.Create(name)
		if err != nil {
			t.Fatalf("zip create: %v", err)
		}
		if _, err := w.Write([]byte(data)); err != nil {
			t.Fatalf("zip write: %v", err)
		}
	}
	if err := zw.Close(); err != nil {
		t.Fatalf("zip close: %v", err)
	}
	return buf.Bytes()
}

func TestSeparateStems(t *testing.T) {
	zipBody := stemsZipFixture(t, map[string]string{
		"htdemucs/vocals.mp3":    "VOX",
		"htdemucs/no_vocals.mp3": "ACC",
	})
	var gotPath, gotTwoStems, gotModel, gotFormat, 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.Fatalf("parse multipart: %v", err)
		}
		gotTwoStems = r.FormValue("two_stems")
		gotModel = r.FormValue("model")
		gotFormat = r.FormValue("format")
		if _, hdr, err := r.FormFile("file"); err == nil {
			gotFilename = hdr.Filename
		} else {
			t.Errorf("file part: %v", err)
		}
		w.Header().Set("Content-Type", "application/zip")
		_, _ = w.Write(zipBody)
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	ss, err := p.StemSeparatorModel("audioutils")
	if err != nil {
		t.Fatalf("StemSeparatorModel: %v", err)
	}
	res, err := ss.SeparateStems(context.Background(),
		audio.StemSeparationRequest{Audio: []byte("SONG"), MIME: "audio/mpeg"},
		audio.WithStemMode("two"), audio.WithStemModel("htdemucs_ft"), audio.WithStemFormat("mp3"))
	if err != nil {
		t.Fatalf("SeparateStems: %v", err)
	}
	if gotPath != "/upstream/audioutils/v1/stems" {
		t.Errorf("path = %q", gotPath)
	}
	if gotTwoStems != "vocals" || gotModel != "htdemucs_ft" || gotFormat != "mp3" || gotFilename != "audio.mp3" {
		t.Errorf("two_stems/model/format/filename = %q/%q/%q/%q", gotTwoStems, gotModel, gotFormat, gotFilename)
	}
	if len(res.Stems) != 2 {
		t.Fatalf("stems = %+v", res.Stems)
	}
	byName := map[string]audio.Stem{}
	for _, s := range res.Stems {
		byName[s.Name] = s
	}
	if v := byName["vocals"]; string(v.Audio) != "VOX" || v.MIME != "audio/mpeg" {
		t.Errorf("vocals = %+v", v)
	}
	if a := byName["no_vocals"]; string(a.Audio) != "ACC" || a.MIME != "audio/mpeg" {
		t.Errorf("no_vocals = %+v", a)
	}
}

func TestSeparateStemsOmitsUnsetFields(t *testing.T) {
	zipBody := stemsZipFixture(t, map[string]string{"vocals.wav": "V"})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := r.ParseMultipartForm(1 << 20); err != nil {
			t.Fatalf("parse multipart: %v", err)
		}
		for _, k := range []string{"two_stems", "model", "format"} {
			if v, ok := r.MultipartForm.Value[k]; ok {
				t.Errorf("unset request sent %q = %v, want omitted", k, v)
			}
		}
		_, _ = w.Write(zipBody)
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	ss, _ := p.StemSeparatorModel("audioutils")
	res, err := ss.SeparateStems(context.Background(),
		audio.StemSeparationRequest{Audio: []byte("SONG")})
	if err != nil {
		t.Fatalf("SeparateStems: %v", err)
	}
	// Top-level entry, wav extension.
	if len(res.Stems) != 1 || res.Stems[0].Name != "vocals" || res.Stems[0].MIME != "audio/wav" {
		t.Errorf("stems = %+v", res.Stems)
	}
}

func TestSeparateStemsRejectsBadArgs(t *testing.T) {
	p := New(WithBaseURL("http://unused"))
	ss, _ := p.StemSeparatorModel("audioutils")
	if _, err := ss.SeparateStems(context.Background(), audio.StemSeparationRequest{}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("no audio: err = %v, want ErrUnsupported", err)
	}
	if _, err := ss.SeparateStems(context.Background(),
		audio.StemSeparationRequest{Audio: []byte{1}, Mode: "three"}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("mode three: err = %v, want ErrUnsupported", err)
	}
	if _, err := ss.SeparateStems(context.Background(),
		audio.StemSeparationRequest{Audio: []byte{1}, Format: "flac"}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("format flac: err = %v, want ErrUnsupported", err)
	}
}

func TestSeparateStemsRejectsNonZip(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()))
	ss, _ := p.StemSeparatorModel("audioutils")
	_, err := ss.SeparateStems(context.Background(), audio.StemSeparationRequest{Audio: []byte("SONG")})
	var apiErr *llm.APIError
	if !errors.As(err, &apiErr) {
		t.Fatalf("err = %v, want APIError for non-zip body", err)
	}
}

func TestSeparateStemsRejectsEmptyZip(t *testing.T) {
	zipBody := stemsZipFixture(t, map[string]string{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/zip")
		_, _ = w.Write(zipBody)
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	ss, _ := p.StemSeparatorModel("audioutils")
	_, err := ss.SeparateStems(context.Background(), audio.StemSeparationRequest{Audio: []byte("SONG")})
	var apiErr *llm.APIError
	if !errors.As(err, &apiErr) {
		t.Fatalf("err = %v, want APIError for stem-less zip", err)
	}
}

// wavFixture is a minimal RIFF/WAVE header so http.DetectContentType sniffs
// audio/wave.
func wavFixture() []byte {
	return []byte("RIFF\x24\x00\x00\x00WAVEfmt ")
}

func TestSFXGenerate(t *testing.T) {
	var gotPath string
	var gotBody map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		_ = json.NewDecoder(r.Body).Decode(&gotBody)
		w.Header().Set("Content-Type", "audio/wav")
		_, _ = w.Write(wavFixture())
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	sm, err := p.SFXModel("sfxgen-stableaudio")
	if err != nil {
		t.Fatalf("SFXModel: %v", err)
	}
	res, err := sm.Generate(context.Background(),
		musicgen.Request{Prompt: "glass shattering", DurationSeconds: 8},
		musicgen.WithSteps(50), musicgen.WithCFGScale(7), musicgen.WithSeed(42))
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if gotPath != "/upstream/sfxgen-stableaudio/v1/sfx" {
		t.Errorf("path = %q", gotPath)
	}
	want := map[string]any{"prompt": "glass shattering", "seconds": 8.0, "steps": 50.0, "cfg_scale": 7.0, "seed": 42.0}
	for k, w := range want {
		if gotBody[k] != w {
			t.Errorf("%s = %v, want %v", k, gotBody[k], w)
		}
	}
	if res.Audio.MIME != "audio/wav" || len(res.Audio.Data) == 0 {
		t.Errorf("audio = %q/%d bytes", res.Audio.MIME, len(res.Audio.Data))
	}
}

func TestSFXOmitsDefaults(t *testing.T) {
	var gotBody map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewDecoder(r.Body).Decode(&gotBody)
		// No Content-Type: the RIFF sniff must still label the clip.
		w.Header()["Content-Type"] = nil
		_, _ = w.Write(wavFixture())
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	sm, _ := p.SFXModel("sfxgen-stableaudio")
	res, err := sm.Generate(context.Background(), musicgen.Request{Prompt: "boom"})
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	for _, k := range []string{"seconds", "steps", "cfg_scale", "seed"} {
		if v, ok := gotBody[k]; ok {
			t.Errorf("unset request sent %q = %v, want omitted", k, v)
		}
	}
	if res.Audio.MIME != "audio/wav" {
		t.Errorf("MIME = %q, want sniffed audio/wav", res.Audio.MIME)
	}
}

func TestSFXRejectsBadArgs(t *testing.T) {
	p := New(WithBaseURL("http://unused"))
	sm, _ := p.SFXModel("sfxgen-stableaudio")
	if _, err := sm.Generate(context.Background(), musicgen.Request{Prompt: "  "}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("empty prompt: err = %v, want ErrUnsupported", err)
	}
	if _, err := sm.Generate(context.Background(),
		musicgen.Request{Prompt: "boom", Lyrics: "la la"}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("lyrics: err = %v, want ErrUnsupported", err)
	}
	if _, err := sm.Generate(context.Background(),
		musicgen.Request{Prompt: "boom", Format: "mp3"}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("format mp3: err = %v, want ErrUnsupported", err)
	}
}

func TestSFXRejectsNonAudioResponse(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"detail":"queue full"}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	sm, _ := p.SFXModel("sfxgen-stableaudio")
	_, err := sm.Generate(context.Background(), musicgen.Request{Prompt: "boom"})
	var apiErr *llm.APIError
	if !errors.As(err, &apiErr) {
		t.Fatalf("err = %v, want APIError for non-audio body", err)
	}
}

func TestEnhance(t *testing.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.Fatalf("parse multipart: %v", err)
		}
		if _, hdr, err := r.FormFile("file"); err == nil {
			gotFilename = hdr.Filename
		} else {
			t.Errorf("file part: %v", err)
		}
		w.Header().Set("Content-Type", "audio/wav")
		_, _ = w.Write(wavFixture())
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	en, err := p.SpeechEnhancerModel("audioutils")
	if err != nil {
		t.Fatalf("SpeechEnhancerModel: %v", err)
	}
	res, err := en.Enhance(context.Background(),
		audio.EnhancementRequest{Audio: []byte("NOISY"), MIME: "audio/ogg"})
	if err != nil {
		t.Fatalf("Enhance: %v", err)
	}
	if gotPath != "/upstream/audioutils/v1/enhance" {
		t.Errorf("path = %q", gotPath)
	}
	if gotFilename != "audio.ogg" {
		t.Errorf("filename = %q", gotFilename)
	}
	if res.MIME != "audio/wav" || len(res.Audio) == 0 {
		t.Errorf("result = %q/%d bytes", res.MIME, len(res.Audio))
	}
}

func TestEnhanceRejectsEmptyAudio(t *testing.T) {
	p := New(WithBaseURL("http://unused"))
	en, _ := p.SpeechEnhancerModel("audioutils")
	if _, err := en.Enhance(context.Background(), audio.EnhancementRequest{}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("err = %v, want ErrUnsupported", err)
	}
}

func TestEnhanceRejectsNonAudioResponse(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>oops</html>"))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	en, _ := p.SpeechEnhancerModel("audioutils")
	_, err := en.Enhance(context.Background(), audio.EnhancementRequest{Audio: []byte("NOISY")})
	var apiErr *llm.APIError
	if !errors.As(err, &apiErr) {
		t.Fatalf("err = %v, want APIError for non-audio body", err)
	}
}