llmgw/internal/control/catalog/environments_test.go

package catalog

import (
	"strings"
	"testing"
)

const (
	commitA = "0123456789abcdef0123456789abcdef01234567"
	commitB = "fedcba9876543210fedcba9876543210fedcba98"
)

func inEnv(name, env, source string, nodes ...CustomNode) Model {
	m := comfyModel(source, &ComfyUI{Workflow: "w.json", Output: OutputMeshes, Environment: env, Nodes: nodes})
	m.Name = name
	return m
}

func TestModelsInOneEnvironmentShareTheirNodesAndWeights(t *testing.T) {
	sf3d := CustomNode{Repo: "https://github.com/Stability-AI/stable-fast-3d", Commit: commitA,
		Requirements: []string{"trimesh", "transformers>=4.42,<5"}}
	tripo := CustomNode{Repo: "https://github.com/flowtyone/ComfyUI-Flowty-TripoSR", Commit: commitB,
		Requirements: []string{"Transformers >= 4.42, <5", "rembg"}}
	flux := comfyModel("/w/flux.safetensors", &ComfyUI{Environment: "shared"})
	cat, err := New([]Model{
		flux,
		inEnv("sf3d", "shared", "/w/sf3d/model.safetensors", sf3d),
		inEnv("triposr", "shared", "/w/tripo/model.ckpt", tripo),
		inEnv("alone", "", "/w/alone.ckpt"),
	})
	if err != nil {
		t.Fatal(err)
	}
	envs := cat.Environments()
	if len(envs) != 1 {
		t.Fatalf("environments %+v", envs)
	}
	e := envs[0]
	if strings.Join(e.Models, ",") != "flux,sf3d,triposr" {
		t.Errorf("models %v", e.Models)
	}
	if len(e.Sources) != 3 || len(e.Nodes) != 2 {
		t.Errorf("sources %v nodes %+v", e.Sources, e.Nodes)
	}
	// transformers once, with the node that named it first; rembg with its own.
	if got := strings.Join(e.Nodes[0].Requirements, " "); got != "trimesh transformers>=4.42,<5" {
		t.Errorf("first node's requirements %q", got)
	}
	if got := strings.Join(e.Nodes[1].Requirements, " "); got != "rembg" {
		t.Errorf("second node's requirements %q", got)
	}
	if got, ok := cat.EnvironmentOf("sf3d"); !ok || got.Name != "shared" {
		t.Errorf("sf3d's environment %+v %v", got, ok)
	}
	if _, ok := cat.EnvironmentOf("alone"); ok {
		t.Error("a model with no environment is in one")
	}
}

func TestAnEnvironmentWhoseModelsDisagreeIsRefused(t *testing.T) {
	node := func(commit string, reqs ...string) CustomNode {
		return CustomNode{Repo: "https://github.com/a/node", Commit: commit, Requirements: reqs}
	}
	for name, c := range map[string]struct {
		models []Model
		want   string
	}{
		"two commits of one node": {[]Model{
			inEnv("a", "e", "/w/a", node(commitA)), inEnv("b", "e", "/w/b", node(commitB))}, "different commits"},
		"two ranges of one package": {[]Model{
			inEnv("a", "e", "/w/a", node(commitA, "numpy<2")),
			inEnv("b", "e", "/w/b", CustomNode{Repo: "https://github.com/b/other", Commit: commitB,
				Requirements: []string{"numpy>=2"}})}, `"a" and "b" ask for numpy`},
		"a bad name": {[]Model{inEnv("a", "Shared Env", "/w/a")}, "comfyui.environment"},
	} {
		if _, err := New(c.models); err == nil || !strings.Contains(err.Error(), c.want) {
			t.Errorf("%s: %v", name, err)
		}
	}
	// The same node at the same commit, the same range written two ways: fine.
	if _, err := New([]Model{
		inEnv("a", "e", "/w/a", node(commitA, "numpy >= 2")),
		inEnv("b", "e", "/w/b", node(commitA, "NumPy>=2")),
	}); err != nil {
		t.Errorf("agreeing models refused: %v", err)
	}
}

func TestALicenceIsSaidSoAMachineCanMatchIt(t *testing.T) {
	yes := true
	ok := License{ID: "tencent-hunyuan-community", Name: "Tencent Hunyuan Community License", Commercial: &yes,
		ExcludedRegions: []string{"EU", "GB", "KR"}}
	m := comfyModel("/w/f.safetensors", nil)
	m.License = &ok
	if _, err := New([]Model{m}); err != nil {
		t.Fatalf("a full licence refused: %v", err)
	}
	for name, c := range map[string]struct {
		l    License
		want string
	}{
		"no name":        {License{ID: "MIT"}, "license.name"},
		"a sentence id":  {License{ID: "the MIT one", Name: "MIT"}, "license.id"},
		"a lowercase eu": {License{Name: "x", ExcludedRegions: []string{"eu"}}, "excluded_regions"},
		"a negative cap": {License{Name: "x", RevenueCapUSD: -1}, "revenue_cap_usd"},
	} {
		m := comfyModel("/w/f.safetensors", nil)
		l := c.l
		m.License = &l
		if _, err := New([]Model{m}); err == nil || !strings.Contains(err.Error(), c.want) {
			t.Errorf("%s: %v", name, err)
		}
	}
}