llmgw/internal/cli/model_presets_test.go

package cli

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
)

// presetAdmin is an admin API holding settings, which records the models
// added and the settings put.
type presetAdmin struct {
	// api is a control plane that adds presets itself; without it, one
	// from before, which answers their paths with a 404.
	api      bool
	byAPI    []string
	mu       sync.Mutex
	settings string
	added    []catalog.Model
	put      []string
}

func (a *presetAdmin) client(t *testing.T) func() *ctlClient {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.mu.Lock()
		defer a.mu.Unlock()
		body, _ := io.ReadAll(r.Body)
		w.Header().Set("Content-Type", "application/json")
		switch {
		case a.api && r.Method == http.MethodPost && strings.HasPrefix(r.URL.Path, "/api/v1/presets/"):
			name := strings.TrimPrefix(r.URL.Path, "/api/v1/presets/")
			a.byAPI = append(a.byAPI, name)
			_ = json.NewEncoder(w).Encode(adminapi.ModelAdded{Added: name})
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/models":
			var m catalog.Model
			_ = json.Unmarshal(body, &m)
			a.added = append(a.added, m)
			_ = json.NewEncoder(w).Encode(map[string]string{"name": m.Name})
		case r.Method == http.MethodGet && r.URL.Path == "/api/v1/settings":
			_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 1, Body: a.settings})
		case r.Method == http.MethodPut && r.URL.Path == "/api/v1/settings":
			var req adminapi.SettingsRequest
			_ = json.Unmarshal(body, &req)
			a.put = append(a.put, req.Body)
			a.settings = req.Body
			_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 2, Body: req.Body})
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(srv.Close)
	return func() *ctlClient { return newCtlClient(strings.TrimPrefix(srv.URL, "http://"), "") }
}

func runModelAdd(t *testing.T, client func() *ctlClient, args ...string) (string, error) {
	t.Helper()
	cmd := newModelAddCmd(client)
	var out bytes.Buffer
	cmd.SetOut(&out)
	cmd.SetErr(&out)
	cmd.SetArgs(append(args, "--no-source-check"))
	err := cmd.ExecuteContext(t.Context())
	return out.String(), err
}

func TestAPresetIsAddedByTheControlPlane(t *testing.T) {
	a := &presetAdmin{api: true}
	out, err := runModelAdd(t, a.client(t), "triposr")
	if err != nil {
		t.Fatalf("%v\n%s", err, out)
	}
	if len(a.byAPI) != 1 || a.byAPI[0] != "triposr" || len(a.added) != 0 || len(a.put) != 0 {
		t.Errorf("by the API %v; sent the entry %v; settings %v", a.byAPI, a.added, a.put)
	}
}

// A control plane from before presets were its own is sent the entry, and
// ComfyUI set beside it, as the release that brought presets did.
func TestAPresetIsAddedByNameAndHasTheNodesInstallComfyUI(t *testing.T) {
	a := &presetAdmin{settings: "router:\n  retries: 2\n"}
	out, err := runModelAdd(t, a.client(t), "triposr")
	if err != nil {
		t.Fatalf("%v\n%s", err, out)
	}
	if len(a.added) != 1 || a.added[0].Name != "triposr" || a.added[0].ComfyUI.Workflow != "triposr" {
		t.Fatalf("added %+v", a.added)
	}
	if len(a.put) != 1 || !strings.Contains(a.put[0], "comfyui: 0.39.2") || !strings.Contains(a.put[0], "retries: 2") {
		t.Fatalf("settings put %q", a.put)
	}
	if !strings.Contains(out, "install ComfyUI 0.39.2") {
		t.Errorf("not said:\n%s", out)
	}
}

func TestAClusterThatNamesItsComfyUIKeepsIt(t *testing.T) {
	for _, settings := range []string{"engines:\n  comfyui: 0.38.0\n", "engines:\n  comfyui: \"\"\n"} {
		a := &presetAdmin{settings: settings}
		out, err := runModelAdd(t, a.client(t), "flux-schnell")
		if err != nil {
			t.Fatalf("%v\n%s", err, out)
		}
		if len(a.put) != 0 {
			t.Errorf("%q: settings changed to %q", settings, a.put)
		}
		if none := strings.Contains(settings, `""`); none != strings.Contains(out, "set to none") {
			t.Errorf("%q: %s", settings, out)
		}
	}
}

func TestAnUnknownPresetNamesTheOnesThereAre(t *testing.T) {
	a := &presetAdmin{}
	_, err := runModelAdd(t, a.client(t), "sdxl")
	if err == nil || !strings.Contains(err.Error(), "flux-schnell, triposr") || len(a.added) != 0 {
		t.Errorf("%v, added %v", err, a.added)
	}
	if _, err := runModelAdd(t, a.client(t), "triposr", "-f", "x.yaml"); err == nil {
		t.Error("a preset and a file were both taken")
	}
}