llmgw/internal/onboard/onboard_test.go

package onboard

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/doctor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/release"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// fakeMachine answers what onboarding asks a machine, and keeps what it is
// given.
type fakeMachine struct {
	mu       sync.Mutex
	hostname string
	arch     string
	noSudo   bool
	report   doctorReport
	files    map[string][]byte
	scripts  []string
	download error
	// mirror is what curl prints for the mirror's address: 000 unreachable.
	mirror string
}

type doctorReport struct {
	Findings []doctor.Finding `json:"findings"`
	Facts    doctor.Facts     `json:"facts"`
}

var putPath = regexp.MustCompile(`cat > (\S+)\.new`)

func (f *fakeMachine) Run(_ context.Context, script string, stdin []byte) ([]byte, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.scripts = append(f.scripts, script)
	switch {
	case script == "hostname && uname -m":
		return []byte(f.hostname + "\n" + f.arch + "\n"), nil
	case script == "sudo -n true":
		if f.noSudo {
			return []byte("sudo: a password is required"), errors.New("exit 1")
		}
		return nil, nil
	case strings.Contains(script, "install -m 0755") && stdin != nil:
		f.files["/usr/local/bin/llmgw"] = stdin
		return nil, nil
	case strings.HasPrefix(script, "sudo llmgw doctor --json"):
		b, _ := json.Marshal(f.report)
		return b, nil
	case putPath.MatchString(script):
		f.files[putPath.FindStringSubmatch(script)[1]] = stdin
		return nil, nil
	case strings.Contains(script, "systemctl restart llmgw-agent"):
		return []byte("active\n"), nil
	case strings.Contains(script, "curl -s -o /dev/null -w"):
		return []byte(f.mirror), nil
	}
	return nil, nil
}

func (f *fakeMachine) ran(s string) bool {
	for _, x := range f.scripts {
		if strings.Contains(x, s) {
			return true
		}
	}
	return false
}

func newFake() *fakeMachine {
	return &fakeMachine{hostname: "gpu01", arch: map[string]string{"amd64": "x86_64", "arm64": "aarch64"}[runtime.GOARCH],
		files: map[string][]byte{}}
}

func runner(t *testing.T, m *fakeMachine, o Options) (*Runner, *[]Event) {
	t.Helper()
	pki := t.TempDir()
	ca, err := tlsconf.NewCA("test")
	if err != nil {
		t.Fatal(err)
	}
	if err := ca.Write(pki); err != nil {
		t.Fatal(err)
	}
	o.Addr = "10.0.0.21"
	o.PKI, o.EtcDir = pki, t.TempDir()
	if o.Servers == nil {
		o.Servers = []string{"https://10.0.0.10:7443"}
	}
	o.Binary = "onboard_test.go" // anything readable
	AgentUnit = "[Unit]\n"
	var events []Event
	var mu sync.Mutex
	return &Runner{Opts: o, Remote: m, Emit: func(e Event) { mu.Lock(); events = append(events, e); mu.Unlock() }}, &events
}

func TestANodeIsAddedFromItsAddressAlone(t *testing.T) {
	m := newFake()
	r, _ := runner(t, m, Options{Version: "v1.2.3-4-gabc"})
	res, err := r.Run(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	if res.Node != "gpu01" {
		t.Errorf("node id %q, want the machine's hostname", res.Node)
	}
	for _, f := range []string{"/etc/llmgw/ca.pem", "/etc/llmgw/agent.pem", "/etc/llmgw/agent-key.pem",
		"/etc/llmgw/agent.yaml", "/etc/systemd/system/llmgw-agent.service", "/usr/local/bin/llmgw"} {
		if len(m.files[f]) == 0 {
			t.Errorf("%s was not written", f)
		}
	}
	yaml := string(m.files["/etc/llmgw/agent.yaml"])
	if !strings.Contains(yaml, `server: "https://10.0.0.10:7443"`) || strings.Contains(yaml, "node_id") {
		t.Errorf("agent.yaml:\n%s", yaml)
	}
	// The key is never readable by others, not even for a moment on the way.
	if !m.ran("umask 077; cat > /etc/llmgw/agent-key.pem.new") || !m.ran("chmod 0600 /etc/llmgw/agent-key.pem.new") {
		t.Error("the key was not written private")
	}
	if !m.ran("systemctl restart llmgw-agent") {
		t.Error("the agent was not started")
	}
}

// releaseFor serves one release for linux/<arch>, its checksums signed with
// a key the test then trusts; tamper changes the binary after signing.
func releaseFor(t *testing.T, version, arch string, body []byte, signed, tamper bool) string {
	t.Helper()
	pub, priv, err := release.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	t.Setenv("LLMGW_RELEASE_KEY", pub)
	asset := release.AssetName(version, "linux", arch)
	sum := sha256.Sum256(body)
	sums := hex.EncodeToString(sum[:]) + "  " + asset + "\n"
	sig, _ := release.Sign(priv, []byte(sums))
	if tamper {
		body = []byte("tampered")
	}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/" + version + "/checksums.txt":
			_, _ = w.Write([]byte(sums))
		case "/" + version + "/" + release.SignatureName:
			if !signed {
				http.NotFound(w, r)
				return
			}
			_, _ = w.Write([]byte(sig))
		case "/" + version + "/" + asset:
			_, _ = w.Write(body)
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(srv.Close)
	return srv.URL
}

func TestAReleaseIsFetchedAndCheckedHereThenGivenToTheNode(t *testing.T) {
	m := newFake()
	base := releaseFor(t, "v1.2.3", hostArch(), []byte("the release"), true, false)
	r, events := runner(t, m, Options{Version: "v1.2.3", Releases: base})
	if _, err := r.Run(t.Context()); err != nil {
		t.Fatal(err)
	}
	if got := string(m.files["/usr/local/bin/llmgw"]); got != "the release" {
		t.Errorf("the node was given %q", got)
	}
	if m.ran("curl") {
		t.Error("the node fetched something itself")
	}
	if !said(*events, Done, "signature and checksum verified") {
		t.Errorf("the check was not reported: %v", *events)
	}
}

func TestAnUnsignedOrAlteredReleaseIsNotGivenToTheNode(t *testing.T) {
	for name, c := range map[string][2]bool{"unsigned": {false, false}, "altered": {true, true}} {
		m := newFake()
		base := releaseFor(t, "v1.2.3", hostArch(), []byte("the release"), c[0], c[1])
		r, events := runner(t, m, Options{Version: "v1.2.3", Releases: base})
		_, _ = r.Run(t.Context())
		if got := string(m.files["/usr/local/bin/llmgw"]); got == "the release" || got == "tampered" {
			t.Errorf("%s: the node was given the release", name)
		}
		if !said(*events, Warn, "could not be fetched and checked") {
			t.Errorf("%s: the refusal was not said", name)
		}
	}
}

func TestANodeGetsThisMachinesCopyWhenTheReleaseCannotBeReached(t *testing.T) {
	m := newFake()
	r, events := runner(t, m, Options{Version: "v1.2.3", Releases: "https://127.0.0.1:1"})
	if _, err := r.Run(t.Context()); err != nil {
		t.Fatal(err)
	}
	if len(m.files["/usr/local/bin/llmgw"]) == 0 {
		t.Error("the binary was not copied")
	}
	if !said(*events, Warn, "could not be fetched") {
		t.Error("the failed download was not mentioned")
	}
}

func TestAnotherArchitectureIsNotGivenThisMachinesBinary(t *testing.T) {
	m := newFake()
	m.arch = map[string]string{"amd64": "aarch64", "arm64": "x86_64"}[runtime.GOARCH]
	r, _ := runner(t, m, Options{Version: "dev"})
	_, err := r.Run(t.Context())
	if err == nil || !strings.Contains(err.Error(), "cannot be copied") {
		t.Errorf("got %v", err)
	}
}

func TestNoSudoIsSaidPlainly(t *testing.T) {
	m := newFake()
	m.noSudo = true
	r, _ := runner(t, m, Options{})
	_, err := r.Run(t.Context())
	if err == nil || !strings.Contains(err.Error(), "sudo without a password") {
		t.Errorf("got %v", err)
	}
	if len(m.files) > 0 {
		t.Error("the machine was changed before finding it could not be finished")
	}
}

func TestADoctorProblemStopsItUnlessAnswered(t *testing.T) {
	fail := doctor.Finding{Level: doctor.Fail, Area: "nvlink", Message: "fabric manager inactive", Fix: "systemctl enable --now nvidia-fabricmanager"}

	m := newFake()
	m.report.Findings = []doctor.Finding{fail}
	r, events := runner(t, m, Options{Version: "dev"})
	if _, err := r.Run(t.Context()); !errors.Is(err, ErrChecks) {
		t.Fatalf("got %v, want ErrChecks", err)
	}
	if !said(*events, Fail, "fabric manager inactive") {
		t.Error("the problem itself was not reported")
	}
	if _, ok := m.files["/etc/llmgw/agent.pem"]; ok {
		t.Error("a certificate was installed on a machine that failed its checks")
	}

	m = newFake()
	m.report.Findings = []doctor.Finding{fail}
	asked := 0
	r, _ = runner(t, m, Options{Version: "dev", AskContinue: func(n int) bool { asked = n; return true }})
	if _, err := r.Run(t.Context()); err != nil {
		t.Fatal(err)
	}
	if asked != 1 {
		t.Errorf("asked about %d problems, want 1", asked)
	}
}

func ibPort(sm string) discovery.Port {
	return discovery.Port{Num: 1, State: "ACTIVE", LinkLayer: "InfiniBand", SM: &discovery.SMInfo{GUID: sm}}
}

func TestAdaptersAreGroupedByTheFabricTheyReport(t *testing.T) {
	hcas := []discovery.HCA{
		{Name: "mlx5_0", Ports: []discovery.Port{ibPort("0xa")}},
		{Name: "mlx5_1", Ports: []discovery.Port{ibPort("0xa")}},
		{Name: "mlx5_4", Ports: []discovery.Port{ibPort("0xb")}},
		{Name: "mlx5_5", Ports: []discovery.Port{{Num: 1, State: "DOWN"}}},
	}
	got := suggestFabrics(hcas)
	if len(got) != 2 || strings.Join(got[0].HCAs, ",") != "mlx5_0,mlx5_1" || strings.Join(got[1].HCAs, ",") != "mlx5_4" {
		t.Errorf("got %+v", got)
	}
}

func TestTheFabricsChosenAreWrittenToTheAgent(t *testing.T) {
	m := newFake()
	m.report.Facts.HCAs = []discovery.HCA{{Name: "mlx5_0", Ports: []discovery.Port{ibPort("0xa")}}}
	r, _ := runner(t, m, Options{Version: "dev", AskFabrics: func(s []Fabric) []Fabric {
		return []Fabric{{Name: "ib-a", HCAs: s[0].HCAs}}
	}})
	if _, err := r.Run(t.Context()); err != nil {
		t.Fatal(err)
	}
	if yaml := string(m.files["/etc/llmgw/agent.yaml"]); !strings.Contains(yaml, `- name: "ib-a"`) || !strings.Contains(yaml, "hcas: [mlx5_0]") {
		t.Errorf("agent.yaml:\n%s", yaml)
	}

	// Asked nothing, a script is told what to add instead.
	m = newFake()
	m.report.Facts.HCAs = []discovery.HCA{{Name: "mlx5_0", Ports: []discovery.Port{ibPort("0xa")}}}
	r, events := runner(t, m, Options{Version: "dev"})
	if _, err := r.Run(t.Context()); err != nil {
		t.Fatal(err)
	}
	if !said(*events, Warn, "no fabric was given") {
		t.Error("adapters up with no fabric given went unmentioned")
	}
}

func TestANodeNamedOtherThanItsHostnameSaysSo(t *testing.T) {
	res := &Result{Node: "gpu-a", Hostname: "ip-10-0-0-21", Servers: []string{"https://a:7443", "https://b:7443"}}
	y := AgentYAML(Options{}, res)
	if !strings.Contains(y, `node_id: "gpu-a"`) || !strings.Contains(y, "servers:\n  - \"https://a:7443\"") {
		t.Errorf("agent.yaml:\n%s", y)
	}
}

func TestParseFabric(t *testing.T) {
	f, err := ParseFabric("ib-a:mlx5_0, mlx5_1")
	if err != nil || f.Name != "ib-a" || strings.Join(f.HCAs, ",") != "mlx5_0,mlx5_1" {
		t.Errorf("got %+v %v", f, err)
	}
	if _, err := ParseFabric(":mlx5_0"); err == nil {
		t.Error("a fabric with no name was accepted")
	}
}

func said(events []Event, level, s string) bool {
	for _, e := range events {
		if e.Level == level && strings.Contains(e.Detail, s) {
			return true
		}
	}
	return false
}

func TestAnEngineDisabledOnANodeIsWrittenAndNotChecked(t *testing.T) {
	m := newFake()
	r, events := runner(t, m, Options{Version: "dev", DisableEngines: []string{"vllm"}, VLLMVersion: "0.30.0"})
	if _, err := r.Run(t.Context()); err != nil {
		t.Fatal(err)
	}
	if yaml := string(m.files["/etc/llmgw/agent.yaml"]); !strings.Contains(yaml, "engines:\n  disabled: [vllm]") {
		t.Errorf("agent.yaml:\n%s", yaml)
	}
	if !m.ran("sudo llmgw doctor --json --vllm none") {
		t.Errorf("doctor was asked about a vLLM this node will never run: %v", m.scripts)
	}
	if !said(*events, Info, "joins with llama.cpp alone") {
		t.Error("the reduced engine set went unsaid")
	}

	// Not disabled: doctor checks the vLLM the control plane installs, not
	// its own default.
	m = newFake()
	r, _ = runner(t, m, Options{Version: "dev", VLLMVersion: "0.25.1"})
	if _, err := r.Run(t.Context()); err != nil {
		t.Fatal(err)
	}
	if !m.ran("sudo llmgw doctor --json --vllm '0.25.1'") && !m.ran("sudo llmgw doctor --json --vllm 0.25.1") {
		t.Errorf("doctor was not told the control plane's vLLM: %v", m.scripts)
	}
}

func TestTheFabricsGivenAreCheckedAsDeclared(t *testing.T) {
	// agent.yaml is written after doctor runs, so the fabrics on the command
	// line are handed to doctor: checked strictly, and nothing else is.
	m := newFake()
	r, _ := runner(t, m, Options{Version: "dev", Fabrics: []Fabric{{Name: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}}}})
	if _, err := r.Run(t.Context()); err != nil {
		t.Fatal(err)
	}
	if !m.ran("--fabric 'ib-a:mlx5_0,mlx5_1'") {
		t.Errorf("doctor was not told the fabric: %v", m.scripts)
	}
}

func TestTheSudoGrantRemovesItselfWithinTheHour(t *testing.T) {
	inner := sudoGrantScript("ilker")
	if !strings.Contains(inner, "--on-active=3600") || !strings.Contains(inner, "sleep 3600; rm -f "+sudoersDrop) {
		t.Errorf("no removal scheduled: %s", inner)
	}
	if !strings.Contains(inner, "visudo -cf") {
		t.Error("written without visudo's check")
	}
	if out, err := exec.Command("sh", "-n", "-c", inner).CombinedOutput(); err != nil {
		t.Errorf("not shell: %v %s\n%s", err, out, inner)
	}
}

func TestANodeThatCannotReachTheMirrorIsNotAdded(t *testing.T) {
	m := newFake()
	m.mirror = "000"
	r, events := runner(t, m, Options{Addr: "10.0.0.9", MirrorBase: "http://172.17.0.1:8081"})
	var res Result
	err := r.check(t.Context(), &res)
	if !errors.Is(err, ErrChecks) {
		t.Fatalf("a node that cannot reach the mirror was let through: %v", err)
	}
	found := false
	for _, e := range *events {
		if e.Level == Fail && strings.Contains(e.Detail, "http://172.17.0.1:8081/") && strings.Contains(e.Fix, "--base") {
			found = true
		}
	}
	if !found {
		t.Errorf("no failure naming the address and the fix: %+v", *events)
	}

	ok := newFake()
	ok.mirror = "404"
	r, _ = runner(t, ok, Options{Addr: "10.0.0.9", MirrorBase: "http://10.0.0.146:8081"})
	if err := r.check(t.Context(), &res); err != nil {
		t.Errorf("a node that reaches the mirror: %v", err)
	}
}