ImaGen/cmd/imagen/compare_test.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"image/png"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// runCompareWithEnv runs the compare subcommand in a writable tmpdir, with
// XDG_CONFIG_HOME pointing somewhere empty so no host imagen.yaml leaks in.
func runCompareWithEnv(t *testing.T, args []string) (stderr, stdout *bytes.Buffer, runDir string, err error) {
	t.Helper()
	tmp := t.TempDir()
	t.Setenv("XDG_CONFIG_HOME", filepath.Join(tmp, "no-config"))
	t.Setenv("HOME", tmp)

	out := filepath.Join(tmp, "compare")
	// stdlib flag parsing requires flags after the leading positional. Append
	// --output at the end so any caller-supplied flags still parse cleanly.
	args = append(args, "--output", out)

	// Capture stdout/stderr via os pipes since runCompare writes directly.
	oldStdout := os.Stdout
	oldStderr := os.Stderr
	rOut, wOut, _ := os.Pipe()
	rErr, wErr, _ := os.Pipe()
	os.Stdout = wOut
	os.Stderr = wErr
	defer func() {
		os.Stdout = oldStdout
		os.Stderr = oldStderr
	}()

	cmdErr := runCompare(context.Background(), args)

	_ = wOut.Close()
	_ = wErr.Close()
	stdout = &bytes.Buffer{}
	stderr = &bytes.Buffer{}
	_, _ = stdout.ReadFrom(rOut)
	_, _ = stderr.ReadFrom(rErr)

	entries, _ := os.ReadDir(out)
	if len(entries) == 1 {
		runDir = filepath.Join(out, entries[0].Name())
	}
	return stderr, stdout, runDir, cmdErr
}

func TestCompareHappyPathWithMockBackends(t *testing.T) {
	// Two mock instances stand in for two different backends. mock ignores
	// cfg so we can reuse the registered type as the instance name and skip
	// writing imagen.yaml entirely.
	stderr, stdout, runDir, err := runCompareWithEnv(t, []string{
		"a cat in a fishbowl",
		"--models", "mock,mock",
		"--size", "64x64",
		"--seed", "42",
	})
	if err != nil {
		t.Fatalf("runCompare: %v\nstderr: %s", err, stderr.String())
	}

	if runDir == "" {
		t.Fatal("expected a run directory under --output")
	}
	// Sidecar JSON
	sidecar := filepath.Join(runDir, "compare.json")
	data, err := os.ReadFile(sidecar)
	if err != nil {
		t.Fatalf("read sidecar: %v", err)
	}
	var body struct {
		Prompt     string `json:"prompt"`
		Successful int    `json:"successful"`
		Total      int    `json:"total"`
		Results    []struct {
			Backend   string `json:"backend"`
			ImagePath string `json:"image_path"`
			Error     string `json:"error"`
		} `json:"results"`
	}
	if err := json.Unmarshal(data, &body); err != nil {
		t.Fatalf("parse sidecar: %v\n%s", err, data)
	}
	if body.Prompt != "a cat in a fishbowl" {
		t.Errorf("prompt = %q", body.Prompt)
	}
	if body.Total != 2 || body.Successful != 2 {
		t.Errorf("counts = %d successful / %d total", body.Successful, body.Total)
	}
	for _, r := range body.Results {
		if r.Error != "" {
			t.Errorf("backend %s errored: %s", r.Backend, r.Error)
		}
		if _, err := os.Stat(r.ImagePath); err != nil {
			t.Errorf("image not on disk for %s: %v", r.Backend, err)
		}
	}

	// Contact sheet path was printed on stdout.
	sheet := strings.TrimSpace(stdout.String())
	if sheet == "" {
		t.Fatal("expected contact sheet path on stdout")
	}
	f, err := os.Open(sheet)
	if err != nil {
		t.Fatalf("open contact sheet: %v", err)
	}
	defer f.Close()
	img, err := png.Decode(f)
	if err != nil {
		t.Fatalf("decode contact sheet PNG: %v", err)
	}
	if w := img.Bounds().Dx(); w < 100 {
		t.Errorf("contact sheet looks empty (width %d)", w)
	}
}

func TestCompareSkipContactSheet(t *testing.T) {
	stderr, stdout, runDir, err := runCompareWithEnv(t, []string{
		"x",
		"--models", "mock",
		"--size", "32x32",
		"--seed", "1",
		"--no-contact-sheet",
	})
	if err != nil {
		t.Fatalf("runCompare: %v\nstderr: %s", err, stderr.String())
	}
	if got := strings.TrimSpace(stdout.String()); got != "" {
		t.Errorf("expected no stdout output (no contact sheet), got %q", got)
	}
	if _, err := os.Stat(filepath.Join(runDir, "contact-sheet.png")); err == nil {
		t.Errorf("contact-sheet.png should not exist with --no-contact-sheet")
	}
}

func TestCompareRecordsBackendErrors(t *testing.T) {
	// One real (mock) + one unknown. Unknown should fail but not abort the
	// run — sidecar records both, contact sheet built from successes only.
	stderr, _, runDir, err := runCompareWithEnv(t, []string{
		"y",
		"--models", "mock,this-instance-does-not-exist",
		"--size", "32x32",
	})
	if err != nil {
		t.Fatalf("runCompare: %v\nstderr: %s", err, stderr.String())
	}
	sidecar := filepath.Join(runDir, "compare.json")
	data, _ := os.ReadFile(sidecar)
	var body struct {
		Successful int `json:"successful"`
		Total      int `json:"total"`
		Results    []struct {
			Backend string `json:"backend"`
			Error   string `json:"error"`
		} `json:"results"`
	}
	if err := json.Unmarshal(data, &body); err != nil {
		t.Fatalf("parse sidecar: %v", err)
	}
	if body.Total != 2 {
		t.Errorf("expected 2 results, got %d", body.Total)
	}
	if body.Successful != 1 {
		t.Errorf("expected 1 success, got %d", body.Successful)
	}
	var sawError bool
	for _, r := range body.Results {
		if r.Backend == "this-instance-does-not-exist" && r.Error != "" {
			sawError = true
		}
	}
	if !sawError {
		t.Errorf("expected an error recorded for the unknown backend")
	}
}

func TestCompareNoModelsFails(t *testing.T) {
	_, _, _, err := runCompareWithEnv(t, []string{"x"})
	if err == nil {
		t.Fatal("expected error when --models is empty")
	}
	if !strings.Contains(err.Error(), "--models") {
		t.Errorf("error should mention --models, got %v", err)
	}
}

func TestCompareNoPromptFails(t *testing.T) {
	_, _, _, err := runCompareWithEnv(t, []string{"--models", "mock"})
	if err == nil {
		t.Fatal("expected error when prompt is missing")
	}
	if !strings.Contains(err.Error(), "missing prompt") {
		t.Errorf("error should mention missing prompt, got %v", err)
	}
}