Files
llmgw/internal/control/catalog/comfyui_test.go
T
Ilker ManapandClaude Opus 5.5 88781cfcfe A job says how far along it is; a ComfyUI model brings its custom nodes
G2: the gate keeps a websocket open to its ComfyUI under a client id of its
own, puts that id on every workflow it lets through — ComfyUI tells only
the client that queued a workflow how it goes — and keeps the last word on
each: the node running, its step of how many. It answers at
/progress/{prompt_id}. While a job's request waits, the router reads the
queue (queued, n ahead) and the gate (running, node, step, percent) and
writes it into the job, which a client reads back. Only into a job of the
same key: a tag naming another's job writes nothing.

G3, the rest of the plumbing: comfyui.nodes names custom nodes by https
repository and full commit, with the PyPI packages they need. The gate puts
each in the replica's own custom_nodes at that commit and installs the
packages beside the replica, constrained to the versions ComfyUI's
environment already has, which it never changes. The catalog refuses a
branch for a commit and anything but package specs for requirements. An
OBJ a node saves leaves the gateway as a GLB, colours kept, so every mesh
can carry its licence.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-10 18:46:44 +02:00

84 lines
3.5 KiB
Go

package catalog
import (
"strings"
"testing"
)
func comfyModel(source string, c *ComfyUI, profiles ...Profile) Model {
if len(profiles) == 0 {
profiles = []Profile{{Name: "gpu", Engine: EngineComfyUI, GPUClass: "6000-24g", TP: 1}}
}
return Model{Name: "flux", Source: source, ComfyUI: c, Profiles: profiles}
}
func TestAComfyUIModelMakesPicturesFromItsSourcesFile(t *testing.T) {
cat, err := New([]Model{comfyModel("hf:Comfy-Org/flux1-schnell/flux1-schnell-fp8.safetensors", nil)})
if err != nil {
t.Fatal(err)
}
m, _ := cat.Resolve("flux")
if !m.MakesImages() || !m.ServesComfyUI() {
t.Error("a ComfyUI model does not say it makes pictures")
}
if w, ckpt := m.ComfyWorkflow(); w != "" || ckpt != "flux1-schnell-fp8.safetensors" {
t.Errorf("workflow %q, checkpoint %q", w, ckpt)
}
}
func TestAComfyUIModelIsRefusedWhatItCannotDo(t *testing.T) {
for name, c := range map[string]struct {
m Model
want string
}{
"a set without a checkpoint": {comfyModel("https://mirror/flux-dev/v1/", nil), "comfyui.checkpoint"},
"an unknown workflow": {comfyModel("/w/f.safetensors", &ComfyUI{Workflow: "sdxl"}), "no built-in workflow"},
"mixed engines": {comfyModel("/w/f.safetensors", nil,
Profile{Name: "a", Engine: EngineComfyUI, GPUClass: "6000-24g", TP: 1},
Profile{Name: "b", Engine: "vllm", GPUClass: "6000-24g", TP: 1}), "every profile"},
"comfyui settings without the engine": {Model{Name: "x", Source: "/w/x", ComfyUI: &ComfyUI{},
Profiles: []Profile{{Name: "a", Engine: "vllm", GPUClass: "6000-24g", TP: 1}}}, "engine comfyui"},
"no card": {comfyModel("/w/f.safetensors", nil, Profile{Name: "a", Engine: EngineComfyUI, TP: 1}), "gpu_class"},
} {
if _, err := New([]Model{c.m}); err == nil || !strings.Contains(err.Error(), c.want) {
t.Errorf("%s: %v", name, err)
}
}
ok := comfyModel("https://mirror/flux-dev/v1/", &ComfyUI{Workflow: "flux-dev", Checkpoint: "flux1-dev-fp8.safetensors"})
if _, err := New([]Model{ok}); err != nil {
t.Errorf("a set with its checkpoint named was refused: %v", err)
}
}
func TestACustomNodeIsPinnedAndItsPackagesArePackages(t *testing.T) {
node := func(n CustomNode) Model {
return comfyModel("hf:stabilityai/TripoSR/model.ckpt", &ComfyUI{Workflow: "w.json", Output: OutputMeshes,
Nodes: []CustomNode{n}})
}
ok := CustomNode{Repo: "https://github.com/flowtyone/ComfyUI-Flowty-TripoSR",
Commit: "0123456789abcdef0123456789abcdef01234567", Requirements: []string{"omegaconf", "trimesh>=4,<5", "scikit-image"}}
cat, err := New([]Model{node(ok)})
if err != nil {
t.Fatal(err)
}
m, _ := cat.Resolve("flux")
p, _ := m.EngineProfile("gpu")
if !strings.Contains(p.Env[NodesEnv], "ComfyUI-Flowty-TripoSR") || !m.MakesMeshes() || m.MakesImages() {
t.Errorf("env %v meshes %v pictures %v", p.Env, m.MakesMeshes(), m.MakesImages())
}
for name, n := range map[string]CustomNode{
"a branch": {Repo: ok.Repo, Commit: "main"},
"http": {Repo: "http://example.com/n", Commit: ok.Commit},
"an option": {Repo: ok.Repo, Commit: ok.Commit, Requirements: []string{"--index-url=https://evil"}},
"a url": {Repo: ok.Repo, Commit: ok.Commit, Requirements: []string{"x @ https://evil/x.whl"}},
} {
if _, err := New([]Model{node(n)}); err == nil {
t.Errorf("%s: accepted", name)
}
}
noWorkflow := comfyModel("hf:a/b/c.ckpt", &ComfyUI{Output: OutputMeshes})
if _, err := New([]Model{noWorkflow}); err == nil || !strings.Contains(err.Error(), "names its workflow") {
t.Errorf("a mesh model without a workflow: %v", err)
}
}