llmgw/internal/comfy/progress_test.go

package comfy

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestTheBookReadsComfyUIsWebsocket(t *testing.T) {
	b := newProgressBook()
	for _, m := range []string{
		`{"type":"status","data":{"status":{"exec_info":{"queue_remaining":1}}}}`,
		`{"type":"execution_start","data":{"prompt_id":"p1"}}`,
		`{"type":"executing","data":{"node":"5","prompt_id":"p1"}}`,
		`{"type":"progress","data":{"value":12,"max":30,"prompt_id":"p1","node":"5"}}`,
	} {
		b.read([]byte(m))
	}
	p, _ := b.get("p1")
	if p.State != "running" || p.Node != "5" || p.Step != 12 || p.Steps != 30 || p.Percent != 40 {
		t.Errorf("%+v", p)
	}
	b.read([]byte(`{"type":"executing","data":{"node":null,"prompt_id":"p1"}}`))
	if p, _ := b.get("p1"); p.State != "done" || p.Percent != 100 {
		t.Errorf("finished: %+v", p)
	}
	b.read([]byte(`{"type":"execution_error","data":{"prompt_id":"p2","node_type":"KSampler","exception_message":"CUDA out of memory"}}`))
	if p, _ := b.get("p2"); p.State != "failed" || p.Error != "KSampler: CUDA out of memory" {
		t.Errorf("failed: %+v", p)
	}
	for i := range keptPrompts + 10 {
		b.update(fmt.Sprint("x", i), func(*Progress) {})
	}
	if _, ok := b.get("p1"); ok || len(b.by) != keptPrompts {
		t.Errorf("the book kept %d", len(b.by))
	}
}

func TestTheGateHearsTheWorkflowsItLetsThrough(t *testing.T) {
	var mu sync.Mutex
	var sent map[string]any
	comfy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/prompt" {
			mu.Lock()
			_ = json.NewDecoder(r.Body).Decode(&sent)
			mu.Unlock()
			fmt.Fprint(w, `{"prompt_id":"p9"}`)
		}
	}))
	defer comfy.Close()
	u, _ := url.Parse(comfy.URL)
	book := newProgressBook()
	gate := httptest.NewServer(newGate(u, "", "m", book, "llmgw-gate-1"))
	defer gate.Close()
	resp, err := http.Post(gate.URL+"/prompt", "application/json", strings.NewReader(`{"prompt":{},"client_id":"router"}`))
	if err != nil {
		t.Fatal(err)
	}
	b, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	if !strings.Contains(string(b), "p9") {
		t.Fatalf("the answer did not come through: %s", b)
	}
	mu.Lock()
	if sent["client_id"] != "llmgw-gate-1" {
		t.Errorf("ComfyUI was sent client %v: its progress would go to nobody listening", sent["client_id"])
	}
	mu.Unlock()
	book.read([]byte(`{"type":"progress","data":{"value":3,"max":4,"prompt_id":"p9","node":"2"}}`))
	resp, _ = http.Get(gate.URL + "/progress/p9")
	var p Progress
	_ = json.NewDecoder(resp.Body).Decode(&p)
	resp.Body.Close()
	if p.State != "running" || p.Percent != 75 {
		t.Errorf("/progress: %+v", p)
	}
	if resp, _ := http.Get(gate.URL + "/progress/unknown"); resp.StatusCode != http.StatusNotFound {
		t.Errorf("an unknown workflow: %d", resp.StatusCode)
	}
}

func TestAWaitingWorkflowSaysHowManyAreAheadAndThenHowFarAlong(t *testing.T) {
	var mu sync.Mutex
	stage := 0
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		s := stage
		mu.Unlock()
		switch r.URL.Path {
		case "/prompt":
			fmt.Fprint(w, `{"prompt_id":"me"}`)
		case "/queue":
			if s == 0 {
				fmt.Fprint(w, `{"queue_running":[[1,"a"]],"queue_pending":[[3,"me"],[2,"b"],[4,"c"]]}`)
			} else {
				fmt.Fprint(w, `{"queue_running":[[3,"me"]],"queue_pending":[]}`)
			}
		case "/progress/me":
			fmt.Fprint(w, `{"state":"running","node":"5","step":6,"steps":12,"percent":50}`)
		case "/history/me":
			if s < 2 {
				fmt.Fprint(w, `{}`)
				return
			}
			fmt.Fprint(w, `{"me":{"status":{"status_str":"success","completed":true},"outputs":{"9":{"images":[{"filename":"a.png","type":"temp"}]}}}}`)
		case "/view":
			w.Header().Set("Content-Type", "image/png")
			fmt.Fprint(w, "png")
		}
	}))
	defer srv.Close()
	var seen []Progress
	c := &Client{Poll: time.Millisecond}
	done := make(chan error, 1)
	go func() {
		_, err := c.GenerateWatch(t.Context(), srv.URL, "", map[string]any{"5": map[string]any{"class_type": "KSampler"}},
			func(p Progress) {
				mu.Lock()
				seen = append(seen, p)
				stage++
				mu.Unlock()
			})
		done <- err
	}()
	if err := <-done; err != nil {
		t.Fatal(err)
	}
	if len(seen) < 2 || seen[0].State != "queued" || seen[0].Ahead != 2 {
		t.Fatalf("%+v: queued behind the one running and the one numbered before it, wanted first", seen)
	}
	if seen[1].State != "running" || seen[1].Percent != 50 || seen[1].NodeClass != "KSampler" {
		t.Errorf("then %+v", seen[1])
	}
}

func TestANodeIsPutAtItsCommit(t *testing.T) {
	if _, err := exec.LookPath("git"); err != nil {
		t.Skip("no git")
	}
	repo := t.TempDir()
	run := func(args ...string) string {
		cmd := exec.Command("git", args...)
		cmd.Dir = repo
		cmd.Env = append(os.Environ(), "GIT_AUTHOR_NAME=t", "GIT_AUTHOR_EMAIL=t@t", "GIT_COMMITTER_NAME=t", "GIT_COMMITTER_EMAIL=t@t")
		b, err := cmd.CombinedOutput()
		if err != nil {
			t.Fatalf("git %v: %v %s", args, err, b)
		}
		return strings.TrimSpace(string(b))
	}
	run("init", "-q")
	run("config", "uploadpack.allowReachableSHA1InWant", "true")
	_ = os.WriteFile(filepath.Join(repo, "__init__.py"), []byte("v1\n"), 0o644)
	run("add", ".")
	run("commit", "-qm", "one")
	first := run("rev-parse", "HEAD")
	_ = os.WriteFile(filepath.Join(repo, "__init__.py"), []byte("v2\n"), 0o644)
	run("commit", "-qam", "two")

	state := t.TempDir()
	site, err := installNodes(t.Context(), state, "python3", []Node{{Repo: "file://" + repo, Commit: first}}, "", func(string) {})
	if err != nil {
		t.Fatal(err)
	}
	b, _ := os.ReadFile(filepath.Join(state, "custom_nodes", filepath.Base(repo), "__init__.py"))
	if string(b) != "v1\n" || site != "" {
		t.Errorf("the node is %q, site %q: wanted the commit named, not the branch's head", b, site)
	}
	if _, err := installNodes(t.Context(), state, "python3", []Node{{Repo: "file://" + repo, Commit: first,
		Requirements: []string{"trimesh"}}}, "", func(string) {}); err == nil || !strings.Contains(err.Error(), "no uv") {
		t.Errorf("packages without a uv: %v", err)
	}
}