majordomo/provider/llamaswap/mesh_diarize_test.go

package llamaswap

import (
	"context"
	"encoding/base64"
	"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/meshgen"
)

func TestMeshGenerate(t *testing.T) {
	png := pngFixture(t)
	stl := []byte("solid mort\nendsolid mort\n")
	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", "application/octet-stream")
		_, _ = w.Write(stl)
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	mm, err := p.MeshModel("image3d-hunyuan21")
	if err != nil {
		t.Fatalf("MeshModel: %v", err)
	}
	res, err := mm.Generate(context.Background(),
		meshgen.Request{Image: meshgen.Image{MIME: "image/png", Data: png}},
		meshgen.WithFormat("stl"), meshgen.WithSeed(7))
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if gotPath != "/upstream/image3d-hunyuan21/generate" {
		t.Errorf("path = %q", gotPath)
	}
	if gotBody["type"] != "stl" || gotBody["image"] != base64.StdEncoding.EncodeToString(png) {
		t.Errorf("type/image = %v/(b64 mismatch)", gotBody["type"])
	}
	if gotBody["seed"] != float64(7) {
		t.Errorf("seed = %v", gotBody["seed"])
	}
	if _, ok := gotBody["octree_resolution"]; ok {
		t.Error("octree_resolution sent when unset; want omitted")
	}
	if res.Mesh.Format != "stl" || res.Mesh.MIME != "model/stl" || string(res.Mesh.Data) != string(stl) {
		t.Fatalf("mesh = %+v", res.Mesh)
	}
}

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

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	mm, _ := p.MeshModel("image3d-hunyuan21")
	_, err := mm.Generate(context.Background(),
		meshgen.Request{Image: meshgen.Image{Data: pngFixture(t)}})
	var apiErr *llm.APIError
	if !errors.As(err, &apiErr) {
		t.Fatalf("err = %v, want APIError for JSON success body", err)
	}
}

func TestMeshGenerateRejectsBadArgs(t *testing.T) {
	p := New(WithBaseURL("http://unused"))
	mm, _ := p.MeshModel("image3d-hunyuan21")
	if _, err := mm.Generate(context.Background(), meshgen.Request{}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("no image: err = %v, want ErrUnsupported", err)
	}
	if _, err := mm.Generate(context.Background(),
		meshgen.Request{Image: meshgen.Image{Data: []byte{1}}, Format: "fbx"}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("format fbx: err = %v, want ErrUnsupported", err)
	}
}

func TestDiarize(t *testing.T) {
	var gotPath string
	var gotQuery map[string][]string
	var gotField string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		gotQuery = r.URL.Query()
		if err := r.ParseMultipartForm(1 << 20); err != nil {
			t.Errorf("parse multipart: %v", err)
		}
		if f, _, err := r.FormFile("audio_file"); err == nil {
			gotField = "audio_file"
			f.Close()
		}
		_, _ = w.Write([]byte(`{
			"text": "hello there. general kenobi.",
			"language": "en",
			"segments": [
				{"start": 0.0, "end": 1.2, "text": "hello there.", "speaker": "SPEAKER_00"},
				{"start": 1.4, "end": 3.0, "text": "general kenobi.", "speaker": "SPEAKER_01"}
			]
		}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	dm, err := p.DiarizationModel("whisperx-diarize")
	if err != nil {
		t.Fatalf("DiarizationModel: %v", err)
	}
	res, err := dm.Diarize(context.Background(),
		audio.DiarizationRequest{Audio: []byte("RIFFfake"), MIME: "audio/wav"},
		audio.WithSpeakerBounds(2, 4), audio.WithDiarizationLanguage("en"))
	if err != nil {
		t.Fatalf("Diarize: %v", err)
	}
	if gotPath != "/upstream/whisperx-diarize/asr" {
		t.Errorf("path = %q", gotPath)
	}
	for k, want := range map[string]string{
		"output": "json", "diarize": "true",
		"min_speakers": "2", "max_speakers": "4", "language": "en",
	} {
		if got := gotQuery[k]; len(got) != 1 || got[0] != want {
			t.Errorf("query %s = %v, want %q", k, got, want)
		}
	}
	if gotField != "audio_file" {
		t.Error("audio_file part missing")
	}
	if len(res.Segments) != 2 || res.Segments[1].Speaker != "SPEAKER_01" {
		t.Fatalf("segments = %+v", res.Segments)
	}
	if res.Language != "en" || res.Text == "" {
		t.Errorf("language/text = %q/%q", res.Language, res.Text)
	}
}

func TestDiarizeRejectsBadArgs(t *testing.T) {
	p := New(WithBaseURL("http://unused"))
	dm, _ := p.DiarizationModel("whisperx-diarize")
	if _, err := dm.Diarize(context.Background(), audio.DiarizationRequest{}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("no audio: err = %v, want ErrUnsupported", err)
	}
	if _, err := dm.Diarize(context.Background(),
		audio.DiarizationRequest{Audio: []byte{1}, MinSpeakers: 5, MaxSpeakers: 2}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("bounds 5>2: err = %v, want ErrUnsupported", err)
	}
}

func TestDiarizeEmptyTranscriptIsError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"text": "", "segments": []}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	dm, _ := p.DiarizationModel("whisperx-diarize")
	_, err := dm.Diarize(context.Background(), audio.DiarizationRequest{Audio: []byte{1}})
	var apiErr *llm.APIError
	if !errors.As(err, &apiErr) {
		t.Fatalf("err = %v, want APIError for empty transcript", err)
	}
}