ImaGen/internal/backend/workflow_template_test.go

package backend

import (
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"
)

func TestBundledWorkflowsParseable(t *testing.T) {
	names := BundledWorkflowNames()
	if len(names) == 0 {
		t.Fatal("expected at least one bundled workflow")
	}
	mustHave := []string{"flux1-schnell", "flux2-klein", "sd35-medium"}
	for _, want := range mustHave {
		if !slices.Contains(names, want) {
			t.Errorf("bundled workflows missing %q (have: %v)", want, names)
		}
	}
	// Every bundled template must parse and contain at least one node.
	for _, n := range names {
		wf, err := LoadWorkflowTemplate(n)
		if err != nil {
			t.Errorf("LoadWorkflowTemplate(%q): %v", n, err)
			continue
		}
		if len(wf) == 0 {
			t.Errorf("workflow %q has zero nodes", n)
		}
	}
}

func TestLoadWorkflowFromFilesystem(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "custom.json")
	body := `{"1":{"class_type":"X","inputs":{"v":"${prompt}"}}}`
	if err := os.WriteFile(path, []byte(body), 0o644); err != nil {
		t.Fatalf("write tmp workflow: %v", err)
	}
	wf, err := LoadWorkflowTemplate(path)
	if err != nil {
		t.Fatalf("load from path: %v", err)
	}
	if _, ok := wf["1"]; !ok {
		t.Errorf("custom workflow missing node 1")
	}
}

func TestLoadWorkflowUnknownNameErrors(t *testing.T) {
	_, err := LoadWorkflowTemplate("definitely-not-a-real-workflow")
	if err == nil {
		t.Fatal("expected error for unknown workflow name")
	}
	if !strings.Contains(err.Error(), "not found") {
		t.Errorf("error should say not found, got %v", err)
	}
}

func TestSubstituteWorkflowPreservesTypes(t *testing.T) {
	wf := map[string]any{
		"31": map[string]any{
			"class_type": "KSampler",
			"inputs": map[string]any{
				"seed":  "${seed}",
				"steps": "${steps}",
				"text":  "${prompt}",
				"cfg":   "${cfg}",
			},
		},
	}
	subs := map[string]any{
		"seed":   int64(42),
		"steps":  11,
		"prompt": "a cat",
		"cfg":    4.5,
	}
	used, err := SubstituteWorkflow(wf, subs)
	if err != nil {
		t.Fatalf("Substitute: %v", err)
	}
	if len(used) != 4 {
		t.Errorf("used = %v, want all four", used)
	}
	inputs := wf["31"].(map[string]any)["inputs"].(map[string]any)
	if seed, ok := inputs["seed"].(int64); !ok || seed != 42 {
		t.Errorf("seed = %T %v, want int64 42", inputs["seed"], inputs["seed"])
	}
	if steps, ok := inputs["steps"].(int); !ok || steps != 11 {
		t.Errorf("steps = %T %v, want int 11", inputs["steps"], inputs["steps"])
	}
	if text, ok := inputs["text"].(string); !ok || text != "a cat" {
		t.Errorf("text = %T %v, want string", inputs["text"], inputs["text"])
	}
	if cfg, ok := inputs["cfg"].(float64); !ok || cfg != 4.5 {
		t.Errorf("cfg = %T %v, want float64 4.5", inputs["cfg"], inputs["cfg"])
	}
}

func TestSubstituteWorkflowMissingPlaceholderErrors(t *testing.T) {
	wf := map[string]any{
		"1": map[string]any{"inputs": map[string]any{"v": "${missing}"}},
	}
	_, err := SubstituteWorkflow(wf, map[string]any{})
	if err == nil {
		t.Fatal("expected error for missing placeholder")
	}
	if !strings.Contains(err.Error(), "${missing}") {
		t.Errorf("error should name the placeholder, got %v", err)
	}
}

func TestSubstituteWorkflowOnlyWholeTokens(t *testing.T) {
	// Partial-match strings ("prefix ${prompt} suffix") are NOT substituted —
	// the placeholder must be the whole value so we can preserve types.
	wf := map[string]any{
		"1": map[string]any{"inputs": map[string]any{
			"keep_string":  "stuff with ${prompt} inside",
			"replace_full": "${prompt}",
		}},
	}
	used, err := SubstituteWorkflow(wf, map[string]any{"prompt": "x"})
	if err != nil {
		t.Fatalf("Substitute: %v", err)
	}
	inputs := wf["1"].(map[string]any)["inputs"].(map[string]any)
	if inputs["keep_string"].(string) != "stuff with ${prompt} inside" {
		t.Errorf("partial match should be left alone, got %q", inputs["keep_string"])
	}
	if inputs["replace_full"].(string) != "x" {
		t.Errorf("full-value match should substitute, got %q", inputs["replace_full"])
	}
	if _, ok := used["prompt"]; !ok {
		t.Errorf("used should track keys that fired")
	}
}

func TestFlux1SchnellTemplateMatchesLegacyShape(t *testing.T) {
	// Regression guard against the historical hardcoded workflow: every
	// node ID the old Comfy.buildWorkflow used must still exist in the
	// migrated template.
	wf, err := LoadWorkflowTemplate("flux1-schnell")
	if err != nil {
		t.Fatalf("load flux1-schnell: %v", err)
	}
	legacyNodes := []string{"6", "8", "9", "10", "11", "12", "13", "27", "30", "31"}
	for _, id := range legacyNodes {
		if _, ok := wf[id]; !ok {
			t.Errorf("flux1-schnell template missing node %q (legacy parity)", id)
		}
	}
}