llmgw/internal/cli/serverinit_cluster_test.go

package cli

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

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

func TestTheInitAgentUnitIsTheShippedUnit(t *testing.T) {
	b, err := os.ReadFile("../../deploy/systemd/llmgw-agent.service")
	if err != nil {
		t.Fatal(err)
	}
	if string(b) != agentUnit {
		t.Error("deploy/systemd/llmgw-agent.service and the unit 'server init' writes have drifted")
	}
}

func TestTheStarterUsesTheGPUAndFallsBackToCPUOnThisMachine(t *testing.T) {
	m := starterModel("http://10.0.0.5:8081/qwen2.5-0.5b-instruct/starter/q.gguf", []string{"4090-24g"}, true)
	c, err := catalog.New([]catalog.Model{m})
	if err != nil {
		t.Fatalf("the starter is not a valid catalog entry: %v", err)
	}
	got, _ := c.Resolve(starterName)
	if len(got.Profiles) != 2 {
		t.Fatalf("profiles = %+v", got.Profiles)
	}
	gpu, cpu := got.Profiles[0], got.Profiles[1]
	if gpu.GPUClass != "4090-24g" || gpu.Priority <= cpu.Priority {
		t.Errorf("the GPU profile should come first: %+v", gpu)
	}
	if cpu.GPUClass != "" || cpu.Engine != "llamacpp" {
		t.Errorf("the fallback should be llama.cpp on CPU: %+v", cpu)
	}
	for _, p := range got.Profiles {
		if p.NodeSelector[localLabelKey] != localLabelValue {
			t.Errorf("profile %s is not pinned to this machine", p.Name)
		}
	}
	if got.Scaling.Min != 1 {
		t.Error("the starter should be running, not waiting to be asked for")
	}

	if m := starterModel("http://x/q.gguf", nil, true); len(m.Profiles) != 1 || m.Profiles[0].Name != "cpu" {
		t.Errorf("without a GPU only the CPU profile: %+v", m.Profiles)
	}
}

func TestTheMirrorBlockIsAddedAndReads(t *testing.T) {
	in, out := newTestInit(t)
	if err := in.run(t.Context()); err != nil {
		t.Fatalf("%v\n%s", err, out)
	}
	if err := appendMirrorBlockFor(in.configPath(), "http://"+mirrorAPI, "# the mirror"); err != nil {
		t.Fatal(err)
	}
	cfg, err := config.Load(in.configPath())
	if err != nil {
		t.Fatal(err)
	}
	if cfg.Mirror.API != "http://"+mirrorAPI {
		t.Errorf("mirror.api = %q", cfg.Mirror.API)
	}
	if err := cfg.Validate(); err != nil {
		t.Error(err)
	}
	backups, _ := filepath.Glob(in.configPath() + ".bak-*")
	if len(backups) != 1 {
		t.Errorf("backups = %v", backups)
	}
}

func TestTheMirrorUnitAdvertisesAnAddressNodesCanReach(t *testing.T) {
	in := &serverInit{hosts: []string{"10.0.0.140"}}
	u := in.mirrorUnit()
	if !strings.Contains(u, "--base http://10.0.0.140:8081") {
		t.Errorf("base: %s", u)
	}
	// The management port can delete models; nothing but this machine may
	// reach it, which is what lets it go without a certificate.
	if !strings.Contains(u, "--api-listen 127.0.0.1:8082") {
		t.Errorf("the management API is not on loopback: %s", u)
	}
}

func TestWithNoNodeHereTheStarterGoesToTheFirstNodeThatJoins(t *testing.T) {
	m := starterModel("http://x/q.gguf", nil, false)
	for _, p := range m.Profiles {
		if len(p.NodeSelector) > 0 {
			t.Errorf("profile %s is pinned to a machine that is not a node: %v", p.Name, p.NodeSelector)
		}
	}
}

func TestTheMirrorAPIOnlyAnswersItsControlPlane(t *testing.T) {
	allowed, err := parseAllow([]string{"10.0.0.140", "10.1.0.0/24"})
	if err != nil {
		t.Fatal(err)
	}
	h := onlyFrom(allowed, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {}))
	for addr, want := range map[string]int{
		"10.0.0.140:5000": 200, "10.1.0.7:5000": 200, "127.0.0.1:5000": 200, "[::1]:5000": 200,
		"10.0.0.141:5000": 403, "192.168.1.9:5000": 403,
	} {
		r := httptest.NewRequest("GET", "/api/v1/health", nil)
		r.RemoteAddr = addr
		w := httptest.NewRecorder()
		h.ServeHTTP(w, r)
		if w.Code != want {
			t.Errorf("%s: %d, want %d", addr, w.Code, want)
		}
	}
	if _, err := parseAllow([]string{"not-an-address"}); err == nil {
		t.Error("nonsense was accepted")
	}
}

func TestTheInstalledConfigurationIsHardenedAndLoads(t *testing.T) {
	for _, judge := range []bool{false, true} {
		body := fmt.Sprintf(serverConfigTemplate, "/etc/llmgw", "/etc/llmgw", "/etc/llmgw",
			"/var/lib/llmgw/llmgw.db", "ops@example.com", guardSection(judge))
		path := filepath.Join(t.TempDir(), "server.yaml")
		if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
			t.Fatal(err)
		}
		cfg, err := config.Load(path)
		if err != nil {
			t.Fatalf("judge=%v: %v", judge, err)
		}
		if err := cfg.Validate(); err != nil {
			t.Fatalf("judge=%v: %v", judge, err)
		}
		if len(cfg.Guard.Default) != 1 || cfg.Guard.Default[0] != "baseline" {
			t.Errorf("judge=%v: guard.default %v", judge, cfg.Guard.Default)
		}
		p := cfg.Guard.Policies["baseline"]
		if p.Fail != "closed" {
			t.Errorf("judge=%v: the baseline fails %q", judge, p.Fail)
		}
		if (p.Judge != nil) != judge {
			t.Errorf("judge=%v: the policy's judge is %+v", judge, p.Judge)
		}
		if judge && (p.Judge.Model != guardModelName || p.Judge.Input != "block") {
			t.Errorf("judge %+v", p.Judge)
		}
		if cfg.Auth.MaxKeyTTL != 90*24*time.Hour || cfg.Auth.SignupOpen || cfg.DebugPprof || cfg.AdminRemote {
			t.Errorf("not hardened: ttl %v signup %v pprof %v remote %v",
				cfg.Auth.MaxKeyTTL, cfg.Auth.SignupOpen, cfg.DebugPprof, cfg.AdminRemote)
		}
	}
}

func TestTheGuardModelIsAValidJudge(t *testing.T) {
	m := guardModel("http://10.0.0.5:8081/llama-guard-3-1b/guard/g.gguf")
	c, err := catalog.New([]catalog.Model{m})
	if err != nil {
		t.Fatalf("not a valid catalog entry: %v", err)
	}
	got, _ := c.Resolve(guardModelName)
	if !got.IsGuard() || got.Guard.Format != "llamaguard3" || len(got.Guard.Categories) != 13 || got.Scaling.Min != 2 {
		t.Errorf("%+v", got)
	}
	for _, p := range got.Profiles {
		if p.GPUClass != "" {
			t.Errorf("the judge takes a card: %+v", p)
		}
	}
}

func TestInitSaysAtOnceWhenAModelFitsNowhere(t *testing.T) {
	api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/replicas":
			w.Write([]byte(`[]`))
		case "/api/v1/explain-placement":
			w.Write([]byte(`{"model":"qwen","rejections":[{"profile":"cpu","node":"n1","reason":"has 3 GB of memory with 1 free, needs 2"}]}`))
		default:
			http.NotFound(w, r)
		}
	}))
	defer api.Close()
	dir := t.TempDir()
	token := filepath.Join(dir, "token")
	os.WriteFile(token, []byte(strings.Repeat("t", 40)), 0o600)
	os.WriteFile(filepath.Join(dir, "server.yaml"),
		[]byte("admin_listen: "+strings.TrimPrefix(api.URL, "http://")+"\nadmin_token_file: "+token+"\n"), 0o600)

	defer func(g time.Duration) { unplacedGrace = g }(unplacedGrace)
	unplacedGrace = 0
	in := &serverInit{etcDir: dir, out: io.Discard}
	start := time.Now()
	err := in.waitServing(t.Context(), "qwen")
	if err == nil || !strings.Contains(err.Error(), "needs 2") || !strings.Contains(err.Error(), "explain-placement") {
		t.Fatalf("%v", err)
	}
	if time.Since(start) > 10*time.Second {
		t.Errorf("took %s to say so", time.Since(start))
	}
}

func TestAnOldConfigurationGetsGuardrailsOnce(t *testing.T) {
	body := fmt.Sprintf(serverConfigTemplate, "/etc/llmgw", "/etc/llmgw", "/etc/llmgw",
		"/var/lib/llmgw/llmgw.db", "ops@example.com", "")
	path := filepath.Join(t.TempDir(), "server.yaml")
	if err := os.WriteFile(path, []byte(body), 0o640); err != nil {
		t.Fatal(err)
	}
	added, err := appendGuardIfMissing(path, true)
	if err != nil || !added {
		t.Fatalf("added=%v err=%v, want the guard section added", added, err)
	}
	cfg, err := config.Load(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := cfg.Validate(); err != nil {
		t.Fatal(err)
	}
	if len(cfg.Guard.Default) != 1 || cfg.Guard.Policies["baseline"].Judge == nil {
		t.Errorf("guard = %+v, want the baseline with its judge", cfg.Guard)
	}
	if backups, _ := filepath.Glob(path + ".bak-*"); len(backups) != 1 {
		t.Errorf("backups = %v, want the old file kept", backups)
	}
	if added, err := appendGuardIfMissing(path, true); err != nil || added {
		t.Errorf("second run: added=%v err=%v, want nothing changed", added, err)
	}

	// guard: written by hand, even empty, is a decision and stays.
	if err := os.WriteFile(path, []byte(body+"guard: {}\n"), 0o640); err != nil {
		t.Fatal(err)
	}
	if added, err := appendGuardIfMissing(path, true); err != nil || added {
		t.Errorf("empty guard: added=%v err=%v, want it left alone", added, err)
	}
}