llmgw/internal/agent/engines/probe_test.go

package engines

import (
	"context"
	"os"
	"path/filepath"
	"testing"
)

// script writes an executable that stands in for an engine.
func script(t *testing.T, name, body string) string {
	t.Helper()
	dir := t.TempDir()
	p := filepath.Join(dir, name)
	if err := os.WriteFile(p, []byte("#!/bin/sh\n"+body+"\n"), 0o755); err != nil {
		t.Fatal(err)
	}
	return p
}

func TestAnEngineThatIsNotInstalledIsAbsentFromTheReport(t *testing.T) {
	r := NewRegistry(Config{})
	r.engines = map[string]Engine{
		"vllm": &VLLM{cfg: VLLMConfig{Bin: "/nonexistent/vllm"}},
	}
	// Absent, not present-with-unknown-version. Placement reads the keys, and
	// a node listed for an engine it cannot run is the failure being fixed.
	if got := r.Probe(t.Context()); len(got) != 0 {
		t.Errorf("expected nothing, got %v", got)
	}
}

func TestProbingNeverRunsTheEngine(t *testing.T) {
	// A wrapper that shifts past flags it does not recognise and then starts
	// the server. People put these in front of engines, and the agent's own
	// tests install one. Probing with --version would launch an inference
	// server at agent startup, on every node.
	marker := filepath.Join(t.TempDir(), "ran")
	bin := script(t, "fake-vllm", "touch "+marker+"; sleep 60")

	r := NewRegistry(Config{})
	r.engines = map[string]Engine{"vllm": &VLLM{cfg: VLLMConfig{Bin: bin}}}

	got := r.Probe(t.Context())
	if _, ok := got["vllm"]; !ok {
		t.Error("a binary that is there should be reported present")
	}
	if _, err := os.Stat(marker); err == nil {
		t.Fatal("the engine was executed by a probe")
	}
}

func TestAModuleThatWillNotImportIsNotInstalled(t *testing.T) {
	// python3 is on almost every machine, so the interpreter being there says
	// nothing about SGLang. Reporting it present put every node in a fleet
	// forward for an engine none of them had.
	bin := script(t, "python3", "exit 1")
	r := NewRegistry(Config{})
	r.engines = map[string]Engine{"sglang": &SGLang{cfg: SGLangConfig{Python: bin}}}

	if got := r.Probe(t.Context()); len(got) != 0 {
		t.Errorf("a failed import means not installed, got %v", got)
	}
}

func TestAModuleThatImportsIsReportedWithItsVersion(t *testing.T) {
	bin := script(t, "python3", `echo "0.4.1"`)
	r := NewRegistry(Config{})
	r.engines = map[string]Engine{"sglang": &SGLang{cfg: SGLangConfig{Python: bin}}}

	got := r.Probe(t.Context())
	if got["sglang"] != "0.4.1" {
		t.Errorf("expected the version, got %v", got)
	}
}

func TestAVersionIsPickedOutOfWhateverIsPrinted(t *testing.T) {
	// Only the module probe reaches this now, and a Python package version is
	// dotted. A bare build number is deliberately not matched: it would turn a
	// stray number in a line of prose into a version.
	for in, want := range map[string]string{
		"0.4.1":                 "0.4.1",
		"sglang 0.4.post1":      "0.4.post1",
		"warning: noise\n0.4.1": "0.4.1",
		"":                      Unknown,
		"no numbers here":       Unknown,
		"just 4123":             Unknown,
	} {
		if got := firstVersion(in); got != want {
			t.Errorf("firstVersion(%q) = %q, want %q", in, got, want)
		}
	}
}

func TestProbingIsBoundedByTheContext(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	bin := script(t, "python3", "sleep 60")
	r := NewRegistry(Config{})
	r.engines = map[string]Engine{"sglang": &SGLang{cfg: SGLangConfig{Python: bin}}}

	// Registration waits on this. An agent that cannot reach the control plane
	// is worse than one reporting an incomplete list.
	if got := r.Probe(ctx); len(got) != 0 {
		t.Errorf("a cancelled probe should report nothing, got %v", got)
	}
}