llmgw/internal/agent/selfdescribe_test.go

package agent

import (
	"context"
	"os"
	"slices"
	"testing"
	"time"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
)

func TestAnAgentConfigNeedsOnlyWhereTheControlPlaneIs(t *testing.T) {
	// Every node carried its own name, its own address and the same two
	// directories written out by hand, so no two agent files were alike and a
	// copy taken from one node registered the next as the first.
	cfg := Config{Server: "127.0.0.1:7443"}
	if err := cfg.normalize(); err != nil {
		t.Fatal(err)
	}
	host, _ := os.Hostname()
	if cfg.NodeID != host || cfg.Hostname != host {
		t.Errorf("node_id %q, hostname %q; expected the machine's name %q",
			cfg.NodeID, cfg.Hostname, host)
	}
	if cfg.MgmtAddr != "127.0.0.1" {
		t.Errorf("mgmt_addr %q, expected the address on the route to the control plane",
			cfg.MgmtAddr)
	}
	if cfg.WeightsCache != DefaultWeightsCache {
		t.Errorf("weights land in %q", cfg.WeightsCache)
	}
	if !slices.Contains(cfg.ModelDirs, cfg.WeightsCache) {
		t.Errorf("the download directory %q is not looked in: %v",
			cfg.WeightsCache, cfg.ModelDirs)
	}
}

func TestWhatIsWrittenStillWins(t *testing.T) {
	cfg := Config{
		Server: "127.0.0.1:7443", NodeID: "gpu-a-01", MgmtAddr: "10.9.9.9",
		ModelDirs: []string{"/lustre/models"}, WeightsCache: "/scratch/w",
	}
	if err := cfg.normalize(); err != nil {
		t.Fatal(err)
	}
	if cfg.NodeID != "gpu-a-01" || cfg.MgmtAddr != "10.9.9.9" {
		t.Errorf("configured identity overridden: %q %q", cfg.NodeID, cfg.MgmtAddr)
	}
	if !slices.Equal(cfg.ModelDirs, []string{"/lustre/models", "/scratch/w"}) {
		t.Errorf("model dirs %v", cfg.ModelDirs)
	}
}

func TestAHeartbeatSaysWhenTheEnginesChanged(t *testing.T) {
	cfg := Config{Server: "127.0.0.1:7443", NodeID: "n1", WeightsCache: t.TempDir()}
	cfg.Superv.WorkDir = t.TempDir()
	a, err := New(cfg, nil, nil)
	if err != nil {
		t.Fatal(err)
	}
	engines := func(m *llmgwv1.AgentMessage) *llmgwv1.Heartbeat { return m.GetHeartbeat() }

	a.mu.Lock()
	a.configuredVers = map[string]string{"llamacpp": "unknown"}
	a.enginesSent = enginesFingerprint(a.configuredVers)
	a.mu.Unlock()
	if hb := engines(a.heartbeat()); hb.GetEnginesChanged() {
		t.Errorf("unchanged engines were resent")
	}

	a.mu.Lock()
	a.configuredVers = map[string]string{"llamacpp": "unknown", "vllm": "unknown"}
	a.mu.Unlock()
	hb := engines(a.heartbeat())
	if !hb.GetEnginesChanged() || len(hb.GetEngineVersions()) != 2 {
		t.Errorf("an installed engine was not reported: %v", hb.GetEngineVersions())
	}
	if engines(a.heartbeat()).GetEnginesChanged() {
		t.Errorf("reported again on the next heartbeat")
	}
}

func TestWhatTheControlPlaneAsksToInstallIsReportedBack(t *testing.T) {
	cfg := Config{Server: "127.0.0.1:7443", NodeID: "n1", WeightsCache: t.TempDir(),
		EnginesDir: t.TempDir()}
	cfg.Superv.WorkDir = t.TempDir()
	a, err := New(cfg, nil, nil)
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	go a.provisionLoop(ctx)

	// An engine this agent cannot install: no network is needed to learn so,
	// and the answer has to reach the control plane rather than a log.
	a.setProvision(&llmgwv1.Provision{Engines: []*llmgwv1.EngineWant{{Engine: "vllm", Version: "0.29.0"}}})

	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		hb := a.heartbeat().GetHeartbeat()
		if hb.GetProvisionChanged() {
			p := hb.GetProvision()
			if len(p) != 1 || p[0].GetState() != "unsupported" {
				t.Fatalf("reported %v", p)
			}
			a.mu.Lock()
			_, reported := a.engineReportLocked()["vllm"]
			a.mu.Unlock()
			if reported {
				t.Error("an engine that was not installed is reported as present")
			}
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("the installation's outcome never reached a heartbeat")
}