llmgw/internal/cli/serverinit_choices_test.go

package cli

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/pgha"
)

func TestEveryChoiceIsInTheConfigurationWrittenAndItLoads(t *testing.T) {
	in := &serverInit{
		language: "tr", signupOpen: true, defaultRPM: 120, defaultTPM: 50000, maxKeyTTL: 30 * 24 * time.Hour,
		media: true, slurm: true, slurmShared: "here", slurmUsers: "llmgw", evictAfterIdle: time.Minute, minLifetime: 90 * time.Second,
		judge: judgeSpec{Fail: "open", Categories: []string{"Violent", "PII"}, Replicas: 3, Budget: 3 * time.Second},
	}
	j := in.judgeFor()
	body := in.tuneConfig(fmt.Sprintf(serverConfigTemplate, "/etc/llmgw", "/etc/llmgw", "/etc/llmgw",
		"/var/lib/llmgw/llmgw.db", "ops@example.com", guardSectionFor(&j)))
	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("%v\n%s", err, body)
	}
	if err := cfg.Validate(); err != nil {
		t.Fatal(err)
	}
	switch {
	case cfg.Landing.Language != "tr" || !cfg.Auth.SignupOpen:
		t.Errorf("language %q signup %v", cfg.Landing.Language, cfg.Auth.SignupOpen)
	case cfg.Auth.DefaultRPM != 120 || cfg.Auth.DefaultTPM != 50000 || cfg.Auth.MaxKeyTTL != 30*24*time.Hour:
		t.Errorf("limits %d %d %s", cfg.Auth.DefaultRPM, cfg.Auth.DefaultTPM, cfg.Auth.MaxKeyTTL)
	case !cfg.Presets.Auto || cfg.Engines["comfyui"] == "":
		t.Errorf("pictures and 3D: presets %v comfyui %q", cfg.Presets.Auto, cfg.Engines["comfyui"])
	case !cfg.Slurm.Enabled || !cfg.Slurm.AllNodes || cfg.Slurm.Controller != localNodeID():
		t.Errorf("slurm %+v", cfg.Slurm)
	case cfg.Slurm.Shared.Path != "/shared" || cfg.Slurm.Shared.Server != localNodeID() || cfg.Slurm.Users != "llmgw":
		t.Errorf("slurm's shared files and users %+v", cfg.Slurm)
	case cfg.OnDemand.EvictAfterIdle != time.Minute || cfg.OnDemand.MinLifetime != 90*time.Second:
		t.Errorf("on demand %+v", cfg.OnDemand)
	}
	base := cfg.Guard.Policies["baseline"]
	if base.Fail != "open" || base.Judge == nil || !slices.Equal(base.Judge.Categories, []string{"Violent", "PII"}) ||
		base.Judge.LatencyBudget != 3*time.Second {
		t.Errorf("the judge as written: fail %q %+v", base.Fail, base.Judge)
	}
	if !strings.Contains(body, "Three replicas run on CPUs") {
		t.Error("the comment says another number of judges than run")
	}
}

func TestNothingChosenIsTheTemplateAsItWas(t *testing.T) {
	body := fmt.Sprintf(serverConfigTemplate, "/e", "/e", "/e", "/s.db", "a@b", guardSection(true))
	if got := (&serverInit{}).tuneConfig(body); got != body {
		t.Error("an install with nothing chosen writes another configuration than before")
	}
	// Each line tuneConfig changes is in the template, or a choice would be
	// lost without a word.
	for _, line := range []string{"  signup_open: false", "  default_rpm: 60", "  default_tpm: 200000",
		"  max_key_ttl: 2160h   # 90 days", "  # language: tr"} {
		if !strings.Contains(serverConfigTemplate, line) {
			t.Errorf("the template has no %q", line)
		}
	}
}

func TestAStarterIsChosenByNameOrByItsFile(t *testing.T) {
	s, err := starterFrom("qwen3-8b")
	if err != nil || s.SizeGB != 5.0 || !strings.HasPrefix(s.Source, "hf:Qwen/Qwen3-8B-GGUF/") {
		t.Errorf("by name: %+v %v", s, err)
	}
	s, err = starterFrom("hf:bartowski/Phi-4-GGUF/Phi-4-Q4_K_M.gguf")
	if err != nil || s.Name != "phi-4-q4_k_m" {
		t.Errorf("by file: %+v %v", s, err)
	}
	if _, err := starterFrom("gpt-5"); err == nil {
		t.Error("a model nobody offers was taken")
	}
	m := starterModel(s, 8192, 0, "http://x/p.gguf", []string{"6000-24g"}, true)
	if m.Name != "phi-4-q4_k_m" || m.ContextLength != 8192 || len(m.Aliases) != 0 ||
		!slices.Equal(m.Profiles[0].Args, []string{"--ctx-size", "32768"}) {
		t.Errorf("entry %+v", m)
	}
	big, _ := starterFrom("qwen3-14b")
	if m := starterModel(big, 16384, 0, "http://x/q.gguf", nil, true); m.Profiles[0].MinMemoryGB != 11 ||
		!slices.Equal(m.Profiles[0].Args, []string{"--ctx-size", "65536"}) {
		t.Errorf("a 9 GB starter asks for %d GB, %v", m.Profiles[0].MinMemoryGB, m.Profiles[0].Args)
	}
}

func TestTheChoicesAreTheirFlagsAgain(t *testing.T) {
	in := &serverInit{admin: "ops@example.com", hosts: []string{"llm.example.com"}, postgres: true,
		mirrorPath: "/data/models", language: "tr", media: true, slurm: true,
		webhooks: []string{"https://hooks.slack.com/services/T/B/x, json:https://a.example.com/in"},
		judge:    judgeSpec{Categories: []string{"Violent", "Jailbreak"}}, maxKeyTTL: 7 * 24 * time.Hour}
	in.starter, _ = starterFrom("qwen3-4b")
	got := strings.Join(in.initFlags(), " ")
	want := "--admin ops@example.com --host llm.example.com --language tr --max-key-ttl 168h --postgres " +
		"--mirror-dir /data/models --starter qwen3-4b --judge-categories Violent,Jailbreak " +
		"--webhook https://hooks.slack.com/services/T/B/x --webhook json:https://a.example.com/in --media --slurm"
	if got != want {
		t.Errorf("flags\n got %s\nwant %s", got, want)
	}
	if got := strings.Join((&serverInit{}).initFlags(), " "); got != "" {
		t.Errorf("nothing chosen is %q", got)
	}
	if w := shellWord("Non-violent Illegal Acts,PII"); w != "'Non-violent Illegal Acts,PII'" {
		t.Errorf("quoted %s", w)
	}
}

func TestChoicesThatCannotBeInstalledAreRefused(t *testing.T) {
	for _, in := range []serverInit{
		{language: "de"},
		{judge: judgeSpec{Fail: "maybe"}},
		{judge: judgeSpec{Categories: []string{"Rudeness"}}},
		{judge: judgeSpec{Replicas: 9}},
		{starterCtx: 100},
		{maxKeyTTL: time.Minute},
		{vip: "10.0.0.160/24"},
		{standbys: true, vip: "10.0.0.160"},
		{slurm: true, noNode: true},
	} {
		if err := in.checkChoices(); err == nil {
			t.Errorf("taken: %+v", in)
		}
	}
	if err := (&serverInit{standbys: true, vip: "10.0.0.160/24"}).checkChoices(); err != nil {
		t.Error(err)
	}
}

func TestWithStandbysTheMirrorAnswersEveryControlPlane(t *testing.T) {
	in := &serverInit{hosts: []string{"10.0.0.5"}, patroniNodes: []pgha.Node{
		{Member: pgha.Member{Name: "db1", IP: "10.0.0.6"}}, {Member: pgha.Member{Name: "db2", IP: "10.0.0.7"}},
	}}
	if u := in.mirrorUnit(); !strings.Contains(u, "--api-listen 127.0.0.1:8082") || strings.Contains(u, "--api-allow") {
		t.Errorf("one control plane: the API is for this machine only:\n%s", u)
	}
	in.standbys = true
	u := in.mirrorUnit()
	for _, want := range []string{"--api-listen :8082", "--api-allow 10.0.0.5", "--api-allow 10.0.0.6", "--api-allow 10.0.0.7"} {
		if !strings.Contains(u, want) {
			t.Errorf("the unit lacks %q:\n%s", want, u)
		}
	}
	if got := in.mirrorAPIURL(); got != "http://10.0.0.5:8082" {
		t.Errorf("the control planes reach the mirror at %s", got)
	}
}

func TestTheFloatingAddressIsInTheCertificate(t *testing.T) {
	in := &serverInit{vip: "10.0.0.160/24"}
	if !slices.Contains(in.certHosts(), "10.0.0.160") {
		t.Errorf("hosts %v", in.certHosts())
	}
}

func TestStandbysWantTheStoresMachines(t *testing.T) {
	in := &serverInit{standbys: true, out: &strings.Builder{}}
	if err := in.ensureStandbys(context.Background()); err == nil || !strings.Contains(err.Error(), "--patroni") {
		t.Errorf("standbys with nowhere to go: %v", err)
	}
}

func TestTwoClustersOnOneNetworkKeepApartFloatingAddresses(t *testing.T) {
	a, b := vrrpRouterID("10.0.0.161"), vrrpRouterID("10.0.0.164")
	if a == b || a < 1 || a > 255 || b < 1 || b > 255 {
		t.Errorf("router ids %d and %d", a, b)
	}
	conf := keepalivedConf("ens18", "10.0.0.164", 24, 100, "secret12")
	if !strings.Contains(conf, fmt.Sprintf("virtual_router_id %d\n", b)) || !strings.Contains(conf, "interface ens18") {
		t.Errorf("conf:\n%s", conf)
	}
}

func TestTheJudgesAndTheStarterShareTheCPUsRatherThanEachTakingAll(t *testing.T) {
	for _, c := range []struct {
		in             serverInit
		judge, starter int
	}{
		// One machine: one judge, on up to four cores; the starter has the rest.
		{serverInit{cpus: 4}, 4, 0},
		{serverInit{cpus: 8}, 4, 4},
		{serverInit{cpus: 32}, 4, 8},
		{serverInit{cpus: 2}, 2, 0},
		{serverInit{cpus: 1}, 0, 0},
		{serverInit{cpus: 4, noGuardModel: true}, judgeCores, 4}, // no judge: all for the starter
		// Two asked for on one machine: two cores each, one each where two
		// would not fit both.
		{serverInit{cpus: 8, judge: judgeSpec{Replicas: 2}}, 2, 4},
		{serverInit{cpus: 2, judge: judgeSpec{Replicas: 2}}, 1, 0},
		{serverInit{cpus: 8, judge: judgeSpec{Replicas: 3}}, 2, 2},
		{serverInit{cpus: 4, noNode: true}, 2, 0}, // on nodes not known yet
	} {
		got := c.in.cpuPlan()
		if got.judge != c.judge || got.starter != c.starter {
			t.Errorf("%+v: judge %d starter %d, want %d %d", c.in, got.judge, got.starter, c.judge, c.starter)
		}
	}
	m := guardModel("http://x/g.gguf", 2, 1, nil)
	if p := m.Profiles[0]; p.CPUCores != 1 || p.MinMemoryGB != 2 || !slices.Contains(p.Args, "--cache-ram") {
		t.Errorf("judge profile: %d cores, %d GB, args %v", p.CPUCores, p.MinMemoryGB, p.Args)
	}
	s := starterModel(starterChoices[0], defaultContext, 2, "http://x/q.gguf", []string{"6000-24g"}, true)
	if cpu := s.Profiles[len(s.Profiles)-1]; cpu.Name != "cpu" || cpu.CPUCores != 2 || s.Profiles[0].CPUCores != 0 {
		t.Errorf("starter: GPU profile %d cores, CPU %+v", s.Profiles[0].CPUCores, cpu)
	}
}

// A second judge on the same machine shares its cores and is lost with it:
// one machine has one, more machines two.
func TestOneMachineHasOneJudgeAndMoreHaveTwo(t *testing.T) {
	for _, c := range []struct {
		in   serverInit
		want int
	}{
		{serverInit{}, 1},
		{serverInit{noNode: true}, 2},
		{serverInit{patroni: []string{"a@db1", "b@db2"}}, 2},
		{serverInit{standbys: true}, 2},
		{serverInit{judge: judgeSpec{Replicas: 3}}, 3},
	} {
		if got := c.in.judgeFor().Replicas; got != c.want {
			t.Errorf("%+v: %d judges, want %d", c.in, got, c.want)
		}
	}
	// The defaults are left out of the command; a choice against them is in.
	ha := serverInit{noNode: true, patroni: []string{"a@db1", "b@db2"}, judge: judgeSpec{Replicas: 2}}
	if f := strings.Join(ha.initFlags(), " "); strings.Contains(f, "--judge-replicas") {
		t.Errorf("two judges on three machines named: %s", f)
	}
	one := serverInit{judge: judgeSpec{Replicas: 2}}
	if f := strings.Join(one.initFlags(), " "); !strings.Contains(f, "--judge-replicas 2") {
		t.Errorf("two judges on one machine not named: %s", f)
	}
}