llmgw/internal/comfy/gate_test.go

package comfy

import (
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestTheGateWantsTheKeyAndLetsThroughOnlyTheRoutersRequests(t *testing.T) {
	var got []string
	comfy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = append(got, r.Method+" "+r.URL.Path+" auth="+r.Header.Get("Authorization"))
		w.WriteHeader(http.StatusOK)
	}))
	defer comfy.Close()
	u, _ := url.Parse(comfy.URL)
	gate := httptest.NewServer(NewGate(u, "k", "flux"))
	defer gate.Close()

	do := func(method, path, key string) (int, string) {
		req, _ := http.NewRequest(method, gate.URL+path, strings.NewReader("{}"))
		if key != "" {
			req.Header.Set("Authorization", "Bearer "+key)
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		defer resp.Body.Close()
		b, _ := io.ReadAll(resp.Body)
		return resp.StatusCode, string(b)
	}
	if code, _ := do("POST", "/prompt", ""); code != http.StatusUnauthorized {
		t.Errorf("no key: %d", code)
	}
	if code, _ := do("POST", "/prompt", "wrong"); code != http.StatusUnauthorized {
		t.Errorf("wrong key: %d", code)
	}
	for _, p := range []string{"/upload/mask", "/api/manager/queue/install", "/userdata/x", "/internal/logs"} {
		if code, _ := do("POST", p, "k"); code != http.StatusNotFound {
			t.Errorf("%s: %d, want it refused", p, code)
		}
	}
	for _, p := range []string{"/prompt", "/history/p1", "/view?filename=a.png", "/queue", "/interrupt", "/upload/image"} {
		if code, _ := do("POST", p, "k"); code != http.StatusOK {
			t.Errorf("%s: %d", p, code)
		}
	}
	if code, body := do("GET", "/v1/models", "k"); code != 200 || !strings.Contains(body, `"id":"flux"`) {
		t.Errorf("/v1/models: %d %s", code, body)
	}
	if code, _ := do("GET", "/health", "k"); code != 200 {
		t.Errorf("/health: %d", code)
	}
	for _, g := range got {
		if strings.Contains(g, "auth=Bearer") {
			t.Errorf("the key reached ComfyUI: %s", g)
		}
	}
	comfy.Close()
	if code, _ := do("GET", "/health", "k"); code != http.StatusServiceUnavailable {
		t.Errorf("/health with ComfyUI down: %d", code)
	}
}

func TestTheWeightsAreShownByTheirOwnNames(t *testing.T) {
	cache := t.TempDir()
	one := filepath.Join(cache, "0123456789abcdef-flux1-schnell-fp8.safetensors")
	_ = os.WriteFile(one, []byte("w"), 0o644)
	models := filepath.Join(t.TempDir(), "models")
	if err := LinkWeights(one, models); err != nil {
		t.Fatal(err)
	}
	if target, err := os.Readlink(filepath.Join(models, "flux1-schnell-fp8.safetensors")); err != nil || target != one {
		t.Fatalf("%q %v", target, err)
	}

	set := t.TempDir()
	for _, n := range []string{"flux1-dev.safetensors", "ae.safetensors", "t5xxl_fp8.safetensors", ".complete"} {
		_ = os.WriteFile(filepath.Join(set, n), []byte("w"), 0o644)
	}
	if err := LinkWeights(set, models); err != nil {
		t.Fatal(err)
	}
	entries, _ := os.ReadDir(models)
	var names []string
	for _, e := range entries {
		names = append(names, e.Name())
	}
	if strings.Join(names, ",") != "ae.safetensors,flux1-dev.safetensors,t5xxl_fp8.safetensors" {
		t.Errorf("shown %v: the old weights' link should be gone and the marker not shown", names)
	}
	if !strings.Contains(extraModelPaths(models), "checkpoints: .") {
		t.Error(extraModelPaths(models))
	}
}

// TestMain lets the test binary stand in for ComfyUI: started by the gate
// as `python main.py --listen … --port …`, it answers /system_stats and
// says what it was given.
func TestMain(m *testing.M) {
	if os.Getenv("LLMGW_FAKE_COMFYUI") == "1" {
		fakeComfyUI()
		return
	}
	os.Exit(m.Run())
}

func fakeComfyUI() {
	var port, base string
	for i, a := range os.Args {
		switch a {
		case "--port":
			port = os.Args[i+1]
		case "--base-directory":
			base = os.Args[i+1]
		}
	}
	_ = os.WriteFile(filepath.Join(base, "args"), []byte(strings.Join(os.Args[1:], " ")+"\nkey="+os.Getenv("LLMGW_ENGINE_KEY")), 0o644)
	http.HandleFunc("/system_stats", func(w http.ResponseWriter, _ *http.Request) {
		// A progress bar's worth of output with every answer: a write that
		// failed would end this process, as it fails ComfyUI's workflow.
		if _, err := os.Stdout.Write([]byte(strings.Repeat("progress ", 20000) + "\n")); err != nil {
			os.Exit(4)
		}
		_, _ = io.WriteString(w, "{}")
	})
	http.HandleFunc("/quit", func(w http.ResponseWriter, _ *http.Request) { os.Exit(3) })
	if os.Getenv("LLMGW_FAKE_COMFYUI_EXIT") == "1" {
		time.AfterFunc(200*time.Millisecond, func() { os.Exit(3) })
	}
	_ = http.ListenAndServe("127.0.0.1:"+port, nil)
}

func TestTheGateRunsComfyUIAndEndsWithIt(t *testing.T) {
	t.Setenv("LLMGW_FAKE_COMFYUI", "1")
	weights := filepath.Join(t.TempDir(), "flux.safetensors")
	_ = os.WriteFile(weights, []byte("w"), 0o644)
	state := t.TempDir()
	l, _ := net.Listen("tcp", "127.0.0.1:0")
	listen := l.Addr().String()
	l.Close()

	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	done := make(chan error, 1)
	go func() {
		done <- RunGate(ctx, GateConfig{Listen: listen, Model: "flux", Weights: weights, State: state,
			Python: os.Args[0], Dir: t.TempDir(), Args: []string{"--disable-smart-memory"}, Output: io.Discard})
	}()
	get := func(path string) int {
		resp, err := http.Get("http://" + listen + path)
		if err != nil {
			return 0
		}
		resp.Body.Close()
		return resp.StatusCode
	}
	deadline := time.Now().Add(10 * time.Second)
	for get("/health") != http.StatusOK {
		if time.Now().After(deadline) {
			t.Fatal("the gate never said ComfyUI was up")
		}
		time.Sleep(20 * time.Millisecond)
	}
	b, _ := os.ReadFile(filepath.Join(state, "args"))
	for _, want := range []string{"main.py --listen 127.0.0.1 --port", "--base-directory " + state,
		"--extra-model-paths-config " + filepath.Join(state, "extra_model_paths.yaml"), "--disable-smart-memory"} {
		if !strings.Contains(string(b), want) {
			t.Errorf("%q not in what ComfyUI was given: %s", want, b)
		}
	}
	if _, err := os.Lstat(filepath.Join(state, "models", "flux.safetensors")); err != nil {
		t.Error("the weights were not shown to ComfyUI:", err)
	}
	if _, err := os.Stat(filepath.Join(state, "custom_nodes")); err != nil {
		t.Error("ComfyUI lists custom_nodes in its base directory first thing, and it is not there:", err)
	}
	if get("/quit") != http.StatusNotFound {
		t.Error("a path the router never asks for was let through")
	}
	cancel()
	select {
	case err := <-done:
		if err != nil {
			t.Errorf("stopped, the gate said %v", err)
		}
	case <-time.After(25 * time.Second):
		t.Fatal("the gate did not stop")
	}
}

func TestComfyUIExitingEndsTheGate(t *testing.T) {
	t.Setenv("LLMGW_FAKE_COMFYUI", "1")
	t.Setenv("LLMGW_FAKE_COMFYUI_EXIT", "1")
	weights := filepath.Join(t.TempDir(), "w.safetensors")
	_ = os.WriteFile(weights, []byte("w"), 0o644)
	err := RunGate(t.Context(), GateConfig{Listen: "127.0.0.1:0", Weights: weights, State: t.TempDir(),
		Python: os.Args[0], Dir: t.TempDir(), Output: io.Discard})
	if err == nil || !strings.Contains(err.Error(), "exit status 3") {
		t.Fatalf("err = %v", err)
	}
}

type brokenWriter struct{}

func (brokenWriter) Write([]byte) (int, error) { return 0, errors.New("broken pipe") }

func TestComfyUIKeepsWritingWhenWhatTheGateWritesToIsGone(t *testing.T) {
	t.Setenv("LLMGW_FAKE_COMFYUI", "1")
	weights := filepath.Join(t.TempDir(), "w.safetensors")
	_ = os.WriteFile(weights, []byte("w"), 0o644)
	l, _ := net.Listen("tcp", "127.0.0.1:0")
	listen := l.Addr().String()
	l.Close()
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	done := make(chan error, 1)
	go func() {
		done <- RunGate(ctx, GateConfig{Listen: listen, Weights: weights, State: t.TempDir(),
			Python: os.Args[0], Dir: t.TempDir(), Output: brokenWriter{}})
	}()
	ok := 0
	deadline := time.Now().Add(10 * time.Second)
	for ok < 20 && time.Now().Before(deadline) {
		resp, err := http.Get("http://" + listen + "/health")
		if err == nil {
			if resp.StatusCode == http.StatusOK {
				ok++
			}
			resp.Body.Close()
		}
		select {
		case err := <-done:
			t.Fatalf("the gate ended: %v", err)
		default:
		}
		time.Sleep(10 * time.Millisecond)
	}
	if ok < 20 {
		t.Fatalf("ComfyUI answered %d times; output it could not hand on stopped it", ok)
	}
}

func TestLLMGWsOwnNodesArePutInEveryReplica(t *testing.T) {
	state := t.TempDir()
	if err := installOwnNodes(state); err != nil {
		t.Fatal(err)
	}
	b, err := os.ReadFile(filepath.Join(state, "custom_nodes", "llmgw_nodes", "__init__.py"))
	if err != nil || !strings.Contains(string(b), "LLMGWPrepareForMesh") {
		t.Fatalf("%v %.80s", err, b)
	}
}