llmgw/internal/agent/agent_test.go

package agent

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"testing"
	"time"

	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
)

// statusRecorder captures what the control plane hears about shards, and what
// it is asked to adopt.
type statusRecorder struct {
	ch chan status

	mu      sync.Mutex
	calls   []adoptCall
	adopted []reconciler.AdoptedShard
	unclaim []string // shard ids to report back as unclaimed
	fetches []fetchReport
}

type status struct {
	replicaID, shardID string
	state              supervisor.State
	err                string
}

func newRecorder() *statusRecorder { return &statusRecorder{ch: make(chan status, 128)} }

func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) {
	select {
	case r.ch <- status{replicaID, shardID, st, errMsg}:
	default:
	}
}

// fetches records download progress, so a test can assert that a fetch is
// visible from the control plane rather than only in the node's log.
func (r *statusRecorder) OnFetchProgress(replicaID, shardID, source string, done, total int64) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.fetches = append(r.fetches, fetchReport{replicaID, shardID, source, done, total})
}

func (r *statusRecorder) fetchReports() []fetchReport {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]fetchReport(nil), r.fetches...)
}

type fetchReport struct {
	replicaID, shardID, source string
	done, total                int64
}

func (r *statusRecorder) Adopt(nodeID string, shards []reconciler.AdoptedShard) []reconciler.AdoptedShard {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.calls = append(r.calls, adoptCall{nodeID: nodeID, shards: shards})
	r.adopted = append(r.adopted, shards...)

	var unclaimed []reconciler.AdoptedShard
	for _, s := range shards {
		if slices.Contains(r.unclaim, s.ShardID) {
			unclaimed = append(unclaimed, s)
		}
	}
	return unclaimed
}

// adoptCall records that the service reached the reconciler at all, which an
// empty shard list used to prevent.
type adoptCall struct {
	nodeID string
	shards []reconciler.AdoptedShard
}

func (r *statusRecorder) adoptCalls() []adoptCall {
	r.mu.Lock()
	defer r.mu.Unlock()
	return slices.Clone(r.calls)
}

func (r *statusRecorder) adoptedShards() []reconciler.AdoptedShard {
	r.mu.Lock()
	defer r.mu.Unlock()
	return slices.Clone(r.adopted)
}

func (r *statusRecorder) markUnclaimed(ids ...string) {
	r.mu.Lock()
	r.unclaim = append(r.unclaim, ids...)
	r.mu.Unlock()
}

func (r *statusRecorder) await(t *testing.T, limit time.Duration, pred func(status) bool) status {
	t.Helper()
	deadline := time.After(limit)
	var seen []status
	for {
		select {
		case s := <-r.ch:
			seen = append(seen, s)
			if pred(s) {
				return s
			}
		case <-deadline:
			t.Fatalf("no matching status within %v; saw %+v", limit, seen)
			return status{}
		}
	}
}

// fakeEngineScript writes a script that stands in for vLLM.
//
// It parses the real --port flag the vLLM adapter generates rather than being
// told out of band, so the adapter's argument construction is exercised too.
func fakeEngineScript(t *testing.T, body string) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), "fake-vllm")
	script := `#!/bin/sh
port=""
while [ $# -gt 0 ]; do
  case "$1" in
    --port) port="$2"; shift 2 ;;
    *) shift ;;
  esac
done
` + body + "\n"
	if err := os.WriteFile(path, []byte(script), 0o755); err != nil {
		t.Fatal(err)
	}
	return path
}

// servingEngine answers /health and /v1/models, which is what readiness needs.
func servingEngine(t *testing.T) string {
	t.Helper()
	py, err := exec.LookPath("python3")
	if err != nil {
		t.Skip("python3 is required for the fake engine")
	}
	return fakeEngineScript(t, `exec `+py+` -c '
import http.server, json, sys
port = int(sys.argv[1])
class H(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == "/health":
            self.send_response(200); self.end_headers(); return
        if self.path == "/v1/models":
            b = json.dumps({"object":"list","data":[{"id":"test-model"}]}).encode()
            self.send_response(200)
            self.send_header("content-type","application/json")
            self.send_header("content-length", str(len(b)))
            self.end_headers(); self.wfile.write(b); return
        self.send_response(404); self.end_headers()
    def log_message(self, *a): pass
http.server.HTTPServer(("127.0.0.1", port), H).serve_forever()
' "$port"`)
}

type harness struct {
	svc *agentsvc.Service
	inv *inventory.Store
	rec *statusRecorder
	srv *httptest.Server
	ag  *Agent
	ln  *dropListener
	// streams counts how many times an agent has opened the Connect stream,
	// which is the only honest way to tell a reconnect from a stale entry.
	streams *atomic.Int64
}

// dropListener tracks accepted connections so a test can sever them.
//
// httptest's CloseClientConnections does not work here: h2c with prior
// knowledge hijacks the connection, and httptest stops tracking hijacked
// connections. A test relying on it would pass without ever reconnecting.
type dropListener struct {
	net.Listener
	mu    sync.Mutex
	conns []net.Conn
}

func (l *dropListener) Accept() (net.Conn, error) {
	c, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	l.mu.Lock()
	l.conns = append(l.conns, c)
	l.mu.Unlock()
	return c, nil
}

// drop severs every connection, as a control plane restart would.
func (l *dropListener) drop() {
	l.mu.Lock()
	conns := l.conns
	l.conns = nil
	l.mu.Unlock()
	for _, c := range conns {
		_ = c.Close()
	}
}

// start wires a real agent to a real service over h2c, so this exercises the
// protocol rather than two structs calling each other.
func start(t *testing.T, cfgFns ...func(*Config)) *harness {
	t.Helper()
	return startWith(t, slog.New(slog.DiscardHandler), cfgFns...)
}

// startLogged is start with logging on, for diagnosing protocol behaviour.
func startLogged(t *testing.T, cfgFns ...func(*Config)) *harness {
	t.Helper()
	return startWith(t, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug})), cfgFns...)
}

func startWith(t *testing.T, log *slog.Logger, cfgFns ...func(*Config)) *harness {
	t.Helper()

	inv := inventory.NewStore()
	rec := newRecorder()
	svc := agentsvc.New(inv, rec, false, log)

	path, handler := svc.Handler()
	mux := http.NewServeMux()
	mux.Handle(path, handler)

	var streams atomic.Int64
	counting := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		streams.Add(1)
		mux.ServeHTTP(w, r)
	})

	raw, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	ln := &dropListener{Listener: raw}
	// connect's bidirectional streams need HTTP/2; h2c provides it without TLS.
	srv := httptest.NewUnstartedServer(h2c.NewHandler(counting, &http2.Server{}))
	srv.Listener = ln
	srv.Start()

	cfg := DefaultConfig()
	cfg.NodeID = "n1"
	cfg.Server = srv.URL
	cfg.MgmtAddr = "127.0.0.1"
	cfg.HeartbeatInterval = 50 * time.Millisecond
	cfg.ReconnectBase = 20 * time.Millisecond
	cfg.ReconnectMax = 100 * time.Millisecond
	cfg.Discovery.SkipGPUs = true
	cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
	cfg.Superv = supervisor.Config{
		WorkDir: t.TempDir(), LogLines: 64,
		StopTimeout: 2 * time.Second, HealthInterval: 20 * time.Millisecond,
	}
	for _, f := range cfgFns {
		f(&cfg)
	}

	client := &http.Client{Transport: &http2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, addr)
		},
	}}

	ag, err := New(cfg, client, log)
	if err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	go func() { _ = ag.Run(ctx) }()

	h := &harness{svc: svc, inv: inv, rec: rec, srv: srv, ag: ag, ln: ln, streams: &streams}
	t.Cleanup(func() {
		cancel()
		ag.Supervisor().StopAll(context.Background())
		srv.Close()
	})
	h.waitRegistered(t)
	return h
}

func (h *harness) waitRegistered(t *testing.T) {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		if _, ok := h.inv.Get("n1"); ok && len(h.svc.ConnectedNodes()) == 1 {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("the agent never registered")
}

// shardCommand builds a single-node vLLM shard on a free port.
//
// The model path is a real file. It used to be "/models/x", which stopped
// working once the agent began checking that weights exist before launching an
// engine — the check doing exactly what it is for, on a test that was relying
// on nobody looking.
func shardCommand(t *testing.T, shardID string) reconciler.ShardCommand {
	t.Helper()
	modelPath := filepath.Join(t.TempDir(), "model.gguf")
	if err := os.WriteFile(modelPath, []byte("weights"), 0o644); err != nil {
		t.Fatal(err)
	}
	return reconciler.ShardCommand{
		ReplicaID: "rep-1", ShardID: shardID, NodeID: "n1",
		Spec: engines.ShardSpec{
			ReplicaID: "rep-1", ShardID: shardID, Rank: 0, NNodes: 1,
			GPUIndices: []int{0}, HTTPPort: freePort(t),
			Profile: engines.Profile{
				Name: "p", Engine: "vllm", TP: 1,
				ModelPath: modelPath, ServedModelName: "test-model",
			},
		},
		StartupTimeout: 15 * time.Second,
	}
}

func freePort(t *testing.T) int {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	return ln.Addr().(*net.TCPAddr).Port
}

func TestRegistrationCarriesInventory(t *testing.T) {
	// The node must actually have the adapters it declares. It used to declare
	// mlx5_0 and mlx5_1 on a node whose sysfs held neither, which passed
	// because nothing compared the two — and that is the gap a real cluster
	// would have found as a replica hanging on its first all-reduce.
	ibRoot := ibTreeWithActivePorts(t, "mlx5_0", "mlx5_1")
	h := start(t, func(c *Config) {
		c.Hostname = "gpu-a-01"
		c.Labels = map[string]string{"rack": "r12"}
		c.Discovery.IBRoot = ibRoot
		c.Fabrics = []inventory.FabricAttachment{
			{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200},
		}
	})

	n, ok := h.inv.Get("n1")
	if !ok {
		t.Fatal("node missing from inventory")
	}
	if n.Hostname != "gpu-a-01" || n.MgmtAddr != "127.0.0.1" {
		t.Errorf("identity wrong: %+v", n)
	}
	if n.Labels["rack"] != "r12" {
		t.Errorf("labels lost: %v", n.Labels)
	}
	// The fabric label cannot be inferred from the wire, so losing it in
	// transit would silently disable multi-node placement.
	if !n.HasFabric("ib-a") {
		t.Fatalf("fabric attachment lost: %+v", n.Fabrics)
	}
	if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) {
		t.Errorf("HCAs=%v", got)
	}
}

func TestHeartbeatKeepsNodeAlive(t *testing.T) {
	h := start(t)

	n, _ := h.inv.Get("n1")
	first := n.LastSeen

	// Waited for rather than slept through: four heartbeat intervals was not
	// always enough on a loaded runner for one to have landed.
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if n, _ = h.inv.Get("n1"); n.LastSeen.After(first) {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Error("heartbeats are not refreshing LastSeen; the node would be declared lost")
}

func TestStartAndStopShardEndToEnd(t *testing.T) {
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = servingEngine(t)
	})

	cmd := shardCommand(t, "shard-a")
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}

	got := h.rec.await(t, 30*time.Second, func(s status) bool {
		return s.shardID == "shard-a" && s.state == supervisor.Ready
	})
	if got.replicaID != "rep-1" {
		t.Errorf("replica id lost in transit: %q", got.replicaID)
	}

	if err := h.svc.StopShard(t.Context(), "n1", "shard-a"); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-a" && s.state == supervisor.Stopped
	})
}

func TestFailedShardCarriesDiagnosis(t *testing.T) {
	// When a model fails to load the engine's output is the only diagnosis, so
	// it has to survive the trip to the control plane.
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `echo 'CUDA out of memory' >&2; exit 1`)
	})

	if err := h.svc.StartShard(t.Context(), shardCommand(t, "shard-bad")); err != nil {
		t.Fatal(err)
	}
	got := h.rec.await(t, 30*time.Second, func(s status) bool {
		return s.shardID == "shard-bad" && s.state == supervisor.Failed
	})
	if got.err == "" {
		t.Error("the failure carried no error message")
	}
}

func TestEngineThatNeverBecomesReady(t *testing.T) {
	// A process that runs but never serves must time out rather than hang the
	// replica forever.
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
	})

	cmd := shardCommand(t, "shard-stuck")
	cmd.StartupTimeout = 500 * time.Millisecond
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 20*time.Second, func(s status) bool {
		return s.shardID == "shard-stuck" && s.state == supervisor.Failed
	})
}

func TestValidationRejectsBeforeLaunch(t *testing.T) {
	// An engine that refuses a profile does so after minutes of loading;
	// catching it at dispatch costs nothing.
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = servingEngine(t)
	})

	cmd := shardCommand(t, "shard-invalid")
	cmd.Spec.Profile.ModelPath = "" // validation requires one

	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	got := h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-invalid" && s.state == supervisor.Failed
	})
	if got.err == "" {
		t.Error("validation failure carried no reason")
	}
}

func TestDispatchToUnconnectedNode(t *testing.T) {
	h := start(t)

	err := h.svc.StartShard(t.Context(), reconciler.ShardCommand{NodeID: "nowhere", ShardID: "x"})
	if err == nil {
		t.Fatal("dispatching to an unconnected node succeeded")
	}
	if !errors.Is(err, agentsvc.ErrNotConnected) {
		t.Errorf("err=%v, expected ErrNotConnected", err)
	}
}

func TestAgentReconnects(t *testing.T) {
	// A control plane restart must not require touching every GPU node.
	h := start(t)

	before := h.streams.Load()
	h.ln.drop()

	// Counting streams, not connected nodes: a stale entry in the connection
	// map would make a node look connected without any reconnect having
	// happened, and an assertion on it would pass vacuously.
	deadline := time.Now().Add(15 * time.Second)
	for time.Now().Before(deadline) {
		if h.streams.Load() > before && len(h.svc.ConnectedNodes()) == 1 {
			if err := h.svc.DrainShard(t.Context(), "n1", "nothing"); err != nil {
				t.Fatalf("the reconnected stream does not accept commands: %v", err)
			}
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("the agent never reconnected (streams %d -> %d)", before, h.streams.Load())
}

func TestIdempotentStart(t *testing.T) {
	// The control plane retries commands; a repeat must not start a second
	// process (plan §5.6).
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 30`)
	})

	cmd := shardCommand(t, "shard-idem")
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-idem" && s.state == supervisor.Starting
	})

	sup := h.ag.Supervisor()
	sh, err := sup.Get("shard-idem")
	if err != nil {
		t.Fatal(err)
	}
	pid := sh.Status().PID

	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	time.Sleep(400 * time.Millisecond)

	sh2, err := sup.Get("shard-idem")
	if err != nil {
		t.Fatal(err)
	}
	if sh2.Status().PID != pid {
		t.Errorf("a second process was started: pid %d then %d", pid, sh2.Status().PID)
	}
}

// TestRunningShardsReportedOnReconnect covers the control plane restart hole
// (plan §6.5): a reconnecting agent has to describe what it is already running,
// or the control plane concludes the cluster is empty and places duplicates
// while the originals keep holding GPUs.
func TestRunningShardsReportedOnReconnect(t *testing.T) {
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
	})

	cmd := shardCommand(t, "shard-live")
	cmd.Pool = "ib-a/h100-80g"
	cmd.Fabric = "ib-a"
	cmd.Spec.HCAs = []string{"mlx5_0"}

	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-live" && s.state == supervisor.Starting
	})

	// Drop the stream, as a control plane restart would.
	h.ln.drop()

	deadline := time.Now().Add(15 * time.Second)
	for time.Now().Before(deadline) {
		if got := h.rec.adoptedShards(); len(got) > 0 {
			a := got[0]
			if a.ShardID != "shard-live" || a.ReplicaID != "rep-1" {
				t.Fatalf("wrong shard reported: %+v", a)
			}
			// Everything the control plane needs to rebuild the replica.
			if a.Model != "test-model" || a.Profile != "p" {
				t.Errorf("catalog identity lost: model=%q profile=%q", a.Model, a.Profile)
			}
			if a.Pool != "ib-a/h100-80g" || a.Fabric != "ib-a" {
				t.Errorf("pool/fabric lost: %q %q", a.Pool, a.Fabric)
			}
			if !slices.Equal(a.HCAs, []string{"mlx5_0"}) {
				t.Errorf("HCAs lost: %v", a.HCAs)
			}
			if a.Endpoint == "" {
				t.Error("rank 0 reported no endpoint; the router would have nowhere to send")
			}
			if len(a.GPUIndices) == 0 {
				t.Error("GPU indices lost; the inventory could not reclaim them")
			}
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("the reconnecting agent never reported its running shard")
}

// TestUnclaimedShardIsStopped covers the other half: a shard whose model left
// the catalog is holding GPUs nothing will ever route to.
func TestUnclaimedShardIsStopped(t *testing.T) {
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
	})

	cmd := shardCommand(t, "shard-orphaned")
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-orphaned" && s.state == supervisor.Starting
	})

	sh, err := h.ag.Supervisor().Get("shard-orphaned")
	if err != nil {
		t.Fatal(err)
	}
	pid := sh.Status().PID

	// The control plane will refuse to claim it on reconnect.
	h.rec.markUnclaimed("shard-orphaned")
	h.ln.drop()

	deadline := time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) {
		if !processAliveForTest(pid) {
			return
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Error("an unclaimed shard kept running and holding its GPUs")
}

// processAliveForTest mirrors the supervisor's liveness check.
func processAliveForTest(pid int) bool {
	return syscall.Kill(pid, 0) == nil
}

func TestAShardWithMissingWeightsFailsBeforeLaunching(t *testing.T) {
	// The commonest setup mistake on a fresh node. Left to the engine it
	// surfaces minutes into a load, in the engine's own words, on a machine
	// nobody is watching — so it is caught here, where the message can say
	// what is actually wrong.
	a := &Agent{cfg: Config{WeightsCache: t.TempDir()}, log: slog.New(slog.DiscardHandler)}
	_, err := a.weights().Ensure(t.Context(), filepath.Join(t.TempDir(), "not-there.gguf"))
	if err == nil {
		t.Fatal("a missing weights file was accepted")
	}
	if !strings.Contains(err.Error(), "not-there.gguf") {
		t.Errorf("the error does not name the file: %v", err)
	}
}

func TestTheCacheDefaultsToTheFirstModelDirectory(t *testing.T) {
	// A node that has said where its models live has already answered this;
	// asking twice is a second thing to get out of step.
	dir := t.TempDir()
	a := &Agent{cfg: Config{ModelDirs: []string{dir, "/elsewhere"}}, log: slog.New(slog.DiscardHandler)}
	if got := a.weights().CacheDir; got != dir {
		t.Errorf("cache = %q, want the first model directory %q", got, dir)
	}
}

// ibTreeWithActivePorts builds a sysfs InfiniBand tree with the named adapters,
// each with one active port, shaped like the real thing.
func ibTreeWithActivePorts(t *testing.T, names ...string) string {
	t.Helper()
	root := t.TempDir()
	for _, name := range names {
		dir := filepath.Join(root, name, "ports", "1")
		if err := os.MkdirAll(filepath.Join(dir, "gids"), 0o755); err != nil {
			t.Fatal(err)
		}
		for file, content := range map[string]string{
			"state":      "4: ACTIVE",
			"phys_state": "5: LinkUp",
			"rate":       "200 Gb/sec (4X HDR)",
			"link_layer": "InfiniBand",
			"lid":        "0x1a",
			"gids/0":     "fec0:0000:0000:000a:0002:c903:00a0:1234",
		} {
			if err := os.WriteFile(filepath.Join(dir, file), []byte(content+"\n"), 0o644); err != nil {
				t.Fatal(err)
			}
		}
	}
	return root
}

func TestAFabricWhoseAdaptersAreAbsentIsNotAdvertised(t *testing.T) {
	// The failure this prevents: a typo, or a card that enumerated differently
	// after a reboot, produced NCCL_IB_HCA naming an adapter that is not there.
	// The replica starts, passes its health checks, and hangs on its first
	// all-reduce — the exact case the fabric invariant exists to prevent,
	// arriving through a door nobody was watching.
	h := start(t, func(c *Config) {
		c.Discovery.IBRoot = ibTreeWithActivePorts(t, "mlx5_0")
		c.Fabrics = []inventory.FabricAttachment{
			{Fabric: "ib-typo", HCAs: []string{"mlx5_9"}, RateGbps: 200},
		}
	})

	n, ok := h.inv.Get("n1")
	if !ok {
		t.Fatal("node missing from inventory")
	}
	if n.HasFabric("ib-typo") {
		t.Errorf("a fabric with no usable adapter was advertised: %+v", n.Fabrics)
	}
}

func TestAFabricRateIsTheSlowestActiveLink(t *testing.T) {
	// A collective runs at the speed of its slowest link, so reporting the
	// fastest would have the scheduler believe a fabric is quicker than it is.
	root := t.TempDir()
	for name, rate := range map[string]string{"mlx5_0": "200", "mlx5_1": "100"} {
		dir := filepath.Join(root, name, "ports", "1")
		if err := os.MkdirAll(filepath.Join(dir, "gids"), 0o755); err != nil {
			t.Fatal(err)
		}
		for file, content := range map[string]string{
			"state": "4: ACTIVE", "phys_state": "5: LinkUp",
			"rate": rate + " Gb/sec (4X HDR)", "link_layer": "InfiniBand",
			"lid": "0x1a", "gids/0": "fec0::1",
		} {
			if err := os.WriteFile(filepath.Join(dir, file), []byte(content+"\n"), 0o644); err != nil {
				t.Fatal(err)
			}
		}
	}
	h := start(t, func(c *Config) {
		c.Discovery.IBRoot = root
		c.Fabrics = []inventory.FabricAttachment{
			{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}},
		}
	})

	n, _ := h.inv.Get("n1")
	for _, f := range n.Fabrics {
		if f.Fabric == "ib-a" && f.RateGbps != 100 {
			t.Errorf("fabric rate = %d, want the slower link's 100", f.RateGbps)
		}
	}
}

func TestAPortGoingDownIsToldToTheControlPlaneOnce(t *testing.T) {
	// Checked only at registration, a port going down went unnoticed until
	// the agent restarted, and the node went on being offered multi-node
	// work that would hang on its first collective.
	root := ibTreeWithActivePorts(t, "mlx5_0")
	h := start(t, func(c *Config) {
		c.Discovery.IBRoot = root
		c.Fabrics = []inventory.FabricAttachment{{Fabric: "ib-a", HCAs: []string{"mlx5_0"}}}
	})
	if n, _ := h.inv.Get("n1"); !n.HasFabric("ib-a") {
		t.Fatalf("the fabric should be advertised at registration: %+v", n.Fabrics)
	}

	hb := h.ag.heartbeat().GetHeartbeat()
	if hb.GetFabricsChanged() || len(hb.GetIbPorts()) != 1 || hb.GetIbPorts()[0].GetState() != "ACTIVE" {
		t.Fatalf("a quiet heartbeat should carry the port and no change: %+v", hb)
	}

	if err := os.WriteFile(filepath.Join(root, "mlx5_0", "ports", "1", "state"), []byte("1: DOWN\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	hb = h.ag.heartbeat().GetHeartbeat()
	if !hb.GetFabricsChanged() || len(hb.GetFabrics()) != 0 {
		t.Errorf("the port going down should withdraw the fabric: changed=%v fabrics=%v",
			hb.GetFabricsChanged(), hb.GetFabrics())
	}
	if hb = h.ag.heartbeat().GetHeartbeat(); hb.GetFabricsChanged() {
		t.Error("the change should be sent once, not on every heartbeat")
	}
}

func TestAShardUsesThisNodesInterfaceAndRoCEIndex(t *testing.T) {
	// The control plane sends one interface name for every node; a node whose
	// management network is on another had NCCL bind nothing. And which RoCE
	// GID entry routes is a property of this node's adapters.
	a := &Agent{cfg: Config{MgmtInterface: "ens3"}}
	a.inv.HCAs = []discovery.HCA{{Name: "mlx5_0", Ports: []discovery.Port{
		{Num: 1, State: "ACTIVE", LinkLayer: discovery.LinkEthernet, RoCEv2GID: 3}}}}
	spec := engines.ShardSpec{MgmtInterface: "eth0", HCAs: []string{"mlx5_0"}}
	a.localizeSpec(&spec)
	if spec.MgmtInterface != "ens3" {
		t.Errorf("interface = %q", spec.MgmtInterface)
	}
	if spec.GIDIndex == nil || *spec.GIDIndex != 3 {
		t.Errorf("GID index = %v", spec.GIDIndex)
	}

	ib := engines.ShardSpec{HCAs: []string{"mlx5_1"}}
	a.inv.HCAs = append(a.inv.HCAs, discovery.HCA{Name: "mlx5_1", Ports: []discovery.Port{
		{Num: 1, State: "ACTIVE", LinkLayer: discovery.LinkInfiniBand, RoCEv2GID: -1}}})
	a.localizeSpec(&ib)
	if ib.GIDIndex != nil {
		t.Error("an InfiniBand shard was given a GID index")
	}
}

func TestTheManagementInterfaceIsFoundNotAssumed(t *testing.T) {
	// It was defaulted to eth0 before it could be looked for, so it never
	// was: a node on ens3 or bond0 had NCCL bootstrap over nothing.
	cfg := DefaultConfig()
	cfg.Server = "http://127.0.0.1:7443"
	cfg.NodeID = "n1"
	cfg.MgmtAddr = "127.0.0.1"
	if err := cfg.normalize(); err != nil {
		t.Fatal(err)
	}
	if cfg.MgmtInterface == "" || cfg.MgmtInterface == "eth0" {
		t.Errorf("interface = %q, want the loopback interface holding 127.0.0.1", cfg.MgmtInterface)
	}
}

func TestEachFabricsSubnetManagersAreAskedOnce(t *testing.T) {
	// Two adapters on one fabric see one master. Its managers are listed
	// through one of them — the same list twice is a query too many — and
	// every port says which master it sees.
	oldSM, oldSMs := querySM, querySMs
	t.Cleanup(func() { querySM, querySMs = oldSM, oldSMs })
	querySM = func(string, int) (*discovery.SMInfo, error) {
		return &discovery.SMInfo{GUID: "0x01", LID: 1, State: "SMINFO_MASTER"}, nil
	}
	lists := 0
	querySMs = func(string, int) ([]discovery.SMRecord, error) {
		lists++
		return []discovery.SMRecord{{GUID: "0x01", State: "master"}, {GUID: "0x02", State: "standby"}}, nil
	}
	h := start(t, func(c *Config) {
		c.Discovery.IBRoot = ibTreeWithActivePorts(t, "mlx5_0", "mlx5_1")
		c.Fabrics = []inventory.FabricAttachment{{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}}}
	})
	hb := h.ag.heartbeat().GetHeartbeat()
	withList := 0
	for _, p := range hb.GetIbPorts() {
		if p.GetSmGuid() != "0x01" {
			t.Errorf("%s sees master %q", p.GetHca(), p.GetSmGuid())
		}
		if len(p.GetSubnetManagers()) > 0 {
			withList++
		}
	}
	if lists != 1 || withList != 1 {
		t.Errorf("listed %d times, carried on %d ports; want once, on one", lists, withList)
	}
}

func TestAFabricScanIsAskedForAndAnswered(t *testing.T) {
	// One node per fabric maps it; the control plane asks and parses.
	old := scanFabric
	t.Cleanup(func() { scanFabric = old })
	var asked string
	scanFabric = func(_ context.Context, hca string, port int) (string, error) {
		asked = fmt.Sprintf("%s/%d", hca, port)
		return "Switch\t36 \"S-00000000000000a1\"\t\t# \"MF0;leaf-a:QM8700/U1\" enhanced port 0 lid 1 lmc 0\n", nil
	}
	h := start(t)
	r, err := h.svc.ScanFabric(t.Context(), "n1", "mlx5_0", 1)
	if err != nil {
		t.Fatal(err)
	}
	if asked != "mlx5_0/1" || !strings.Contains(r.GetOutput(), "leaf-a") || r.GetError() != "" {
		t.Errorf("asked %q, got %+v", asked, r)
	}
}

func TestABandwidthTestIsRunOnTheAdapterAskedFor(t *testing.T) {
	oldRun, oldHave := runWriteBW, haveWriteBW
	t.Cleanup(func() { runWriteBW, haveWriteBW = oldRun, oldHave })
	haveWriteBW = func() bool { return true }
	var got discovery.WriteBW
	runWriteBW = func(_ context.Context, w discovery.WriteBW) (float64, string, error) {
		got = w
		return 388.2, " #bytes  #iterations  BW peak[Gb/sec]  BW average[Gb/sec]\n", nil
	}
	h := start(t, func(c *Config) { c.Discovery.IBRoot = ibTreeWithActivePorts(t, "mlx5_0") })

	r, err := h.svc.FabricTest(t.Context(), "n1", &llmgwv1.FabricTestRequest{
		Role: "client", Hca: "mlx5_0", Port: 1, Peer: "10.0.0.2", TcpPort: 18515, Seconds: 3})
	if err != nil {
		t.Fatal(err)
	}
	if r.GetGbps() != 388.2 || r.GetError() != "" {
		t.Errorf("reply %+v", r)
	}
	if got.HCA != "mlx5_0" || got.Peer != "10.0.0.2" || got.TCPPort != 18515 || got.GIDIndex != -1 {
		t.Errorf("ran %+v", got)
	}

	// An adapter the node does not have is said so, without running anything.
	got = discovery.WriteBW{}
	r, _ = h.svc.FabricTest(t.Context(), "n1", &llmgwv1.FabricTestRequest{Role: "server", Hca: "mlx5_9", Port: 1, Seconds: 1})
	if !strings.Contains(r.GetError(), "no port mlx5_9/1") || got.HCA != "" {
		t.Errorf("reply %+v, ran %+v", r, got)
	}
}