llmgw/internal/cli/tui_setup_test.go

package cli

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"

	tea "github.com/charmbracelet/bubbletea"

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

func press(m tea.Model, keys ...string) tea.Model {
	for _, k := range keys {
		var msg tea.KeyMsg
		switch k {
		case "enter":
			msg = tea.KeyMsg{Type: tea.KeyEnter}
		case "tab":
			msg = tea.KeyMsg{Type: tea.KeyTab}
		case "down":
			msg = tea.KeyMsg{Type: tea.KeyDown}
		case "up":
			msg = tea.KeyMsg{Type: tea.KeyUp}
		case "space":
			msg = tea.KeyMsg{Type: tea.KeySpace, Runes: []rune(" ")}
		case "left":
			msg = tea.KeyMsg{Type: tea.KeyLeft}
		case "esc":
			msg = tea.KeyMsg{Type: tea.KeyEsc}
		case "right":
			msg = tea.KeyMsg{Type: tea.KeyRight}
		default:
			msg = tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(k)}
		}
		m, _ = m.Update(msg)
	}
	return m
}

func TestTheSetupAsksForWhatHighAvailabilityNeedsAndRefusesWhatWillNotDo(t *testing.T) {
	m := tea.Model(newSetupModel("ctl", []string{"10.0.0.5"}))
	m = press(m, "down", "enter")
	if s := m.(setupModel); s.page != pChoices {
		t.Fatalf("high availability did not lead to the choices: page %d", s.page)
	}
	// Three machines, each a control plane, unless told otherwise.
	if v := m.View(); !strings.Contains(v, "Control plane on all three") || !strings.Contains(v, "two standbys") {
		t.Errorf("the standbys are not offered, ticked:\n%s", v)
	}
	// No email: refused, with a reason.
	m = press(m, "up", "enter")
	if s := m.(setupModel); s.page != pChoices || !strings.Contains(s.choices.err, "email") {
		t.Fatalf("went on without an administrator: %+v", s.choices.err)
	}
	// The refusal is a dialog; a key closes it, and the control plane's
	// settings open with the cursor on the email.
	if !strings.Contains(m.View(), "Not yet") {
		t.Error("the refusal was not shown as a dialog")
	}
	m = press(m, "q")
	if s := m.(setupModel); s.choices.dialog == nil || s.choices.dialog.Fields[s.choices.dialog.focus] != s.choices.admin {
		t.Fatal("the email was not given the cursor")
	}
	m = press(m, "a@example.com", "esc", "enter")
	s := m.(setupModel)
	if s.page != pMachines {
		t.Fatalf("page %d: %s", s.page, s.formErr)
	}
	// This machine's own address is not a database machine.
	s.db[0].addr.SetValue("10.0.0.5")
	s.db[0].user.SetValue("ubuntu")
	s.db[0].auth = authAgent
	s.db[1].addr.SetValue("10.0.0.7")
	s.db[1].user.SetValue("ubuntu")
	s.db[1].auth = authPassword
	if err := s.validateMachines(); !strings.Contains(err, "this machine") {
		t.Errorf("this machine as a database machine: %q", err)
	}
	s.db[0].addr.SetValue("10.0.0.6")
	if err := s.validateMachines(); !strings.Contains(err, "password") {
		t.Errorf("a password login with no password: %q", err)
	}
	s.db[1].secret.SetValue("pw")
	if err := s.validateMachines(); err != "" {
		t.Errorf("a valid pair refused: %s", err)
	}
	cmd := s.command()
	for _, want := range []string{"--admin a@example.com", "--patroni ubuntu@10.0.0.6", "--patroni ubuntu@10.0.0.7", "--no-node", "--standbys"} {
		if !strings.Contains(cmd, want) {
			t.Errorf("command lacks %q:\n%s", want, cmd)
		}
	}
	if strings.Contains(cmd, "pw") {
		t.Error("the password is in the command line")
	}
}

func TestChoosingOneMachineHandsOverToInit(t *testing.T) {
	m := press(tea.Model(newSetupModel("ctl", nil)), "enter").(setupModel)
	if !m.single {
		t.Error("one machine did not hand over")
	}
}

func TestTheChecksHoldTheSetupBackOnlyForAFailure(t *testing.T) {
	m := newSetupModel("ctl", nil)
	m.choice, m.page = 1, pCheck
	m.report = preflight.Report{Findings: []preflight.Finding{{Machine: "a", Check: "disk", Level: preflight.Fail, Fix: "more"}}}
	if s := press(m, "enter").(setupModel); s.page != pCheck {
		t.Error("went on past a failed check")
	}
	m.report.Findings[0].Level = preflight.Warn
	if s := press(m, "enter").(setupModel); s.page != pPlan {
		t.Error("a warning held it back")
	}
}

func TestTheConsoleOffersToReplicateOnlyAStoreThatIsNot(t *testing.T) {
	m := consoleModel{}
	if !m.entryOK(cReplicate) {
		t.Error("a single store is not offered replication")
	}
	m.data.sum.Store = []adminapi.StoreMember{{Name: "a", Role: "leader"}}
	if m.entryOK(cReplicate) {
		t.Error("a replicated store is offered replication")
	}
	// A node is asked for in a dialog, then the console closes with it to do.
	var mm tea.Model = consoleModel{menu: cAddNode, loaded: true, width: 100, height: 30}
	mm = press(mm, "enter", "ubuntu@10.0.0.9", "enter")
	if a := mm.(consoleModel).action; a.kind != "node" || a.targets[0] != "ubuntu@10.0.0.9" {
		t.Errorf("%+v", a)
	}
}

func TestTheConsoleWatchesTheClusterAndEditsItsSettings(t *testing.T) {
	dir := t.TempDir()
	file := filepath.Join(dir, "server.yaml")
	if err := os.WriteFile(file, []byte("# mine\nlisten: \":8080\"\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	m := newConsoleModel("127.0.0.1:1", file)
	m.width, m.height, m.loaded = 120, 40, true
	m.data.nodes = []adminapi.NodeView{{ID: "cuda01", State: "ready", GPUTotal: 1, GPUs: []adminapi.GPUView{{Index: 0, Name: "Quadro RTX 6000"}}}}
	m.data.replicas = []adminapi.ReplicaView{{ID: "r1", Model: "coder-14b", State: "ready",
		Shards: []adminapi.ShardView{{Node: "cuda01", GPUs: []int{0}}}}}

	// Watching: the GPUs view says which model holds each card.
	m.menu = cGPUs
	if v := m.View(); !strings.Contains(v, "Quadro RTX 6000") || !strings.Contains(v, "coder-14b") {
		t.Errorf("the GPUs view:\n%s", v)
	}

	// Settings: into router, change retries, save.
	var mm tea.Model = m
	mm.(consoleModel).View()
	cm := mm.(consoleModel)
	cm.menu, cm.inPanel = cSettings, true
	for i, r := range cm.settingsRows() {
		if r.label == "router" {
			cm.row = i
		}
	}
	mm = press(cm, "enter") // into router
	cm = mm.(consoleModel)
	for i, r := range cm.settingsRows() {
		if r.label == "retries" {
			cm.row = i
		}
	}
	if v := cm.View(); !strings.Contains(v, "(default)") {
		t.Errorf("an unset key does not show its default:\n%s", v)
	}
	mm = press(cm, "enter") // the edit dialog
	cm = mm.(consoleModel)
	cm.prompt.SetValue("5")
	mm = press(cm, "enter", "s")
	if d := mm.(consoleModel).dialog; d != "note" {
		t.Fatalf("after saving the dialog is %q: %s", d, mm.(consoleModel).dialogText)
	}
	b, _ := os.ReadFile(file)
	if !strings.Contains(string(b), "retries: 5") || !strings.Contains(string(b), "# mine") {
		t.Errorf("server.yaml:\n%s", b)
	}
}

func TestTheConsoleAddsAControlPlaneWithAFloatingAddress(t *testing.T) {
	replicated := consoleModel{menu: cAddControl, loaded: true, width: 100, height: 30}
	replicated.data.sum.Store = []adminapi.StoreMember{{Name: "a", Role: "leader"}}

	// The address, then the floating address; the console closes to run it.
	var mm tea.Model = replicated
	mm = press(mm, "enter", "ubuntu@10.0.0.150", "enter", "10.0.0.161/24", "enter")
	want := []string{"server", "add", "ubuntu@10.0.0.150", "--vip", "10.0.0.161/24"}
	if a := mm.(consoleModel).action; a.kind != "run" || !slices.Equal(a.targets, want) {
		t.Errorf("%+v", a)
	}

	// No floating address is an answer too.
	mm = press(replicated, "enter", "ubuntu@10.0.0.150", "enter", "enter")
	if a := mm.(consoleModel).action; !slices.Equal(a.targets, want[:3]) {
		t.Errorf("without one: %+v", a)
	}

	// A mistyped one is caught here rather than half way through.
	mm = press(replicated, "enter", "ubuntu@10.0.0.150", "enter", "10.0.0.161", "enter")
	if c := mm.(consoleModel); c.action.kind != "" || !strings.Contains(c.dialogText, "prefix") {
		t.Errorf("an address without its prefix went through: %+v", c.action)
	}

	// On a store on one machine it says what to do first.
	mm = press(consoleModel{menu: cAddControl, loaded: true, width: 100, height: 30}, "enter")
	if c := mm.(consoleModel); c.dialog != "error" || !strings.Contains(c.dialogText, "Replicate the store") {
		t.Errorf("dialog %q: %s", c.dialog, c.dialogText)
	}
}

func TestTheConsoleOnAStandbySaysWhoseClusterItShows(t *testing.T) {
	m := consoleModel{menu: cOverview, loaded: true, width: 120, height: 40}
	m.data.sum.ControlPlanes = adminapi.ControlPlanes{Members: []adminapi.ControlPlane{
		{Name: "llmgw-ctl", State: "leading"},
		{Name: "llmgw-ctl02", State: "standing by", This: true},
	}}
	if v := m.View(); !strings.Contains(v, "Standing by. Shown, and changed, on the leader: llmgw-ctl") {
		t.Errorf("the overview on a standby:\n%s", v)
	}
	m.data.sum.ControlPlanes.Members[0].This, m.data.sum.ControlPlanes.Members[1].This = true, false
	if v := m.View(); strings.Contains(v, "Standing by.") {
		t.Error("the leader's console says it is standing by")
	}
}

func TestTheConsoleShowsTheGuardrails(t *testing.T) {
	m := consoleModel{menu: cGuard, loaded: true, width: 140, height: 60}
	m.data.models = []adminapi.ModelView{{Name: "llama-guard-3-1b", Ready: 2, Desired: 2}}
	m.guard.view = adminapi.GuardView{
		Policies: []string{"baseline"}, Default: []string{"baseline"},
		Details: []adminapi.GuardPolicyView{{Name: "baseline", Mode: "enforce", Fail: "closed", Judge: "llama-guard-3-1b",
			Checks: []string{"input: pii email → redact"}}},
		Attachments: []adminapi.GuardAttachmentView{{Name: "acme/ml"}},
	}
	m.guard.view.Attachments[0].Scope, m.guard.view.Attachments[0].Policy = "team", "baseline"
	m.guard.events = adminapi.GuardEventsView{Counts: []adminapi.GuardCountView{{Action: "block", Policy: "baseline", Category: "S1", Count: 7}}}
	v := m.View()
	for _, want := range []string{"Guardrails", "baseline", "fail closed", "judge llama-guard-3-1b (2/2 ready)",
		"input: pii email → redact", "acme/ml", "IN THE LAST 24H", "S1"} {
		if !strings.Contains(v, want) {
			t.Errorf("%q missing:\n%s", want, v)
		}
	}
	// W moves the window, and asks again.
	m.inPanel = true
	mm, cmd := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("w")})
	if c := mm.(consoleModel); guardWindows[c.guard.window] != "7d" || cmd == nil {
		t.Errorf("window %q", guardWindows[c.guard.window])
	}
}

func TestTheConsoleEditsTheClustersSettingsInTheStore(t *testing.T) {
	dir := t.TempDir()
	file := filepath.Join(dir, "server.yaml")
	if err := os.WriteFile(file, []byte("listen: \":8080\"\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	var put adminapi.SettingsRequest
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == "GET" && r.URL.Path == "/api/v1/settings":
			_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 7, Body: "router:\n  retries: 2\n"})
		case r.Method == "PUT" && r.URL.Path == "/api/v1/settings":
			_ = json.NewDecoder(r.Body).Decode(&put)
			_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 8, Body: put.Body})
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()
	m := newConsoleModel(strings.TrimPrefix(srv.URL, "http://"), file)
	m.width, m.height, m.loaded = 120, 40, true
	m.menu, m.inPanel = cSettings, true
	m.loadSettings()
	for i, r := range m.settingsRows() {
		if r.label == "router" {
			m.row = i
		}
	}
	var mm tea.Model = press(m, "enter")
	cm := mm.(consoleModel)
	for i, r := range cm.settingsRows() {
		if r.label == "retries" {
			cm.row = i
		}
	}
	if v := cm.View(); !strings.Contains(v, "2") {
		t.Errorf("the cluster's value is not shown:\n%s", v)
	}
	mm = press(cm, "enter")
	cm = mm.(consoleModel)
	cm.prompt.SetValue("5")
	mm = press(cm, "enter", "s")
	if d := mm.(consoleModel).dialog; d != "note" || !strings.Contains(mm.(consoleModel).dialogText, "version 8") {
		t.Fatalf("after saving: %q %s", d, mm.(consoleModel).dialogText)
	}
	if put.Base != 7 || !strings.Contains(put.Body, "retries: 5") {
		t.Errorf("sent %+v", put)
	}
	if b, _ := os.ReadFile(file); string(b) != "listen: \":8080\"\n" {
		t.Errorf("the cluster's key went into the machine's file:\n%s", b)
	}
}

func TestTheConsoleWipesOnlyOnTheMachinesName(t *testing.T) {
	m := newConsoleModel("127.0.0.1:1", filepath.Join(t.TempDir(), "server.yaml"))
	m.width, m.height, m.loaded = 120, 40, true
	m.menu = cWipe
	m.dialog, m.dialogTitle, m.dialogText = "wipe", "Wipe this machine", m.wipeQuestion()
	m.prompt = newTextField()
	m.prompt.Focus()
	if v := m.View(); !strings.Contains(v, "Type the machine's name") {
		t.Errorf("the dialog does not ask for the name:\n%s", v)
	}
	m.prompt.SetValue("not-this-one")
	mm := press(m, "enter")
	if a := mm.(consoleModel).action; a.kind != "" {
		t.Fatalf("a wrong name wiped: %+v", a)
	}
	cm := mm.(consoleModel)
	host, _ := os.Hostname()
	cm.prompt.SetValue(host)
	mm = press(cm, "enter")
	if a := mm.(consoleModel).action; a.kind != "wipe" {
		t.Errorf("the right name did not: %+v", a)
	}
}