llmgw/internal/agent/provision/python_test.go

package provision

import (
	"context"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"testing"

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

var gpuNode = discovery.Host{OSID: "ubuntu", OSVersion: "24.04", Arch: "amd64",
	NVIDIADriver: "580.65.06", MinComputeCapability: 8.0}

// fakePython is uv and the environments it builds.
type fakePython struct {
	fakeSystem
	pmu      sync.Mutex
	commands []string
	failPip  bool
}

func (f *fakePython) run(ctx context.Context, env []string, name string, args ...string) ([]byte, error) {
	line := name + " " + strings.Join(args, " ")
	f.pmu.Lock()
	f.commands = append(f.commands, line)
	f.pmu.Unlock()
	switch {
	case strings.HasSuffix(name, "/uv") && len(args) == 1 && args[0] == "--version":
		return []byte("uv 0.12.20\n"), nil
	case strings.HasSuffix(name, "/uv") && args[0] == "venv":
		return nil, os.MkdirAll(filepath.Join(args[len(args)-1], "bin"), 0o755)
	case strings.HasSuffix(name, "/uv") && args[0] == "pip" && args[1] == "install":
		if f.failPip {
			return []byte("error: Hash mismatch for torch\n"), errExit
		}
		venv := filepath.Dir(filepath.Dir(args[slices.Index(args, "--python")+1]))
		return nil, os.WriteFile(filepath.Join(venv, "bin", "vllm"), []byte("#!/bin/sh\n"), 0o755)
	case strings.HasSuffix(name, "/bin/vllm") && args[0] == "--version":
		return []byte("0.30.0\n"), nil
	}
	return f.fakeSystem.run(ctx, env, name, args...)
}

func (f *fakePython) runStdin(_ context.Context, _ []string, stdin, name string, args ...string) ([]byte, error) {
	f.pmu.Lock()
	f.commands = append(f.commands, name+" "+strings.Join(args, " ")+" <<< "+strings.TrimSpace(stdin))
	f.pmu.Unlock()
	out := args[slices.Index(args, "--output-file")+1]
	return nil, os.WriteFile(out, []byte("vllm==0.30.0 --hash=sha256:ab\n"), 0o644)
}

func pythonInstaller(t *testing.T, host discovery.Host, sys *fakePython) (*Installer, *release) {
	uv := newRelease(map[string][]byte{
		"uv-x86_64-unknown-linux-gnu.tar.gz": tarGz(t,
			map[string]string{"uv-x86_64-unknown-linux-gnu/uv": "#!/bin/sh\n"}, nil),
	})
	gh := uv.serve(t)
	mux := http.NewServeMux()
	mux.HandleFunc("/pypi/vllm/0.30.0/json", func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(`{"urls":[{"upload_time_iso_8601":"2026-09-22T05:22:21.647189Z"},
			{"upload_time_iso_8601":"2026-09-22T05:23:10.039243Z"}]}`))
	})
	pypi := httptest.NewServer(mux)
	t.Cleanup(pypi.Close)
	return &Installer{
		Root: t.TempDir(), Host: host, Run: sys.run, RunStdin: sys.runStdin,
		APIBase: gh.URL, DownloadBase: gh.URL, PyPIBase: pypi.URL, HTTP: gh.Client(),
		FreeSpace: func(string) (float64, error) { return 100, nil },
	}, uv
}

func TestVLLMIsInstalledOnAGPUNode(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	in, _ := pythonInstaller(t, gpuNode, sys)

	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
	if got[0].State != Ready {
		t.Fatalf("not ready: %+v", got[0])
	}
	dir := filepath.Join(in.Root, "vllm", "0.30.0-cu13")
	if got[0].Binary != filepath.Join(dir, "bin", "vllm") {
		t.Errorf("binary %s, want the environment's vllm", got[0].Binary)
	}
	if !sys.installed["gcc"] {
		t.Error("gcc, which Triton compiles its launcher with, was not installed")
	}
	all := strings.Join(sys.commands, "\n")
	for _, want := range []string{
		// PyPI as it stood a day after vLLM's last file, so the dependencies
		// are the ones it was released against.
		"--exclude-newer 2026-09-23T05:23:10Z",
		"--generate-hashes",
		"--require-hashes",
		"venv --relocatable --python 3.12",
	} {
		if !strings.Contains(all, want) {
			t.Errorf("no %q in:\n%s", want, all)
		}
	}
	if _, err := os.Stat(filepath.Join(dir, lockName)); err != nil {
		t.Errorf("the resolved requirements were not kept: %v", err)
	}
	if _, err := os.Stat(filepath.Join(in.Root, ".uv-cache")); err == nil {
		t.Error("the download cache was left behind after a successful install")
	}

	// Checked again on the next pass, not installed again.
	before := len(sys.commands)
	if again := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}}); again[0].State != Ready {
		t.Fatalf("second pass: %+v", again[0])
	}
	for _, c := range sys.commands[before:] {
		if strings.Contains(c, "pip install") {
			t.Errorf("a ready environment was installed again: %s", c)
		}
	}
}

func TestVLLMIsNotForEveryNode(t *testing.T) {
	cases := []struct {
		name  string
		host  discovery.Host
		state string
		says  string
	}{
		{"no GPU", ubuntu, Skipped, "has none"},
		{"old driver", discovery.Host{Arch: "amd64", NVIDIADriver: "550.54.15"}, Unsupported, "driver 580"},
		{"V100", discovery.Host{Arch: "amd64", NVIDIADriver: "580.65.06", MinComputeCapability: 7.0}, Unsupported, "7.5"},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
			in, uv := pythonInstaller(t, c.host, sys)
			got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
			if got[0].State != c.state || !strings.Contains(got[0].Detail, c.says) {
				t.Errorf("got %s %q, want %s saying %q", got[0].State, got[0].Detail, c.state, c.says)
			}
			if uv.gets != 0 || len(sys.commands) != 0 {
				t.Errorf("a node that cannot run it still started installing: %v", sys.commands)
			}
		})
	}
}

func TestAPythonEngineIsRefusedWithoutTheDiskForIt(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{"gcc": true}}}
	in, uv := pythonInstaller(t, gpuNode, sys)
	in.FreeSpace = func(string) (float64, error) { return 12.5, nil }
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "12.5 GB") {
		t.Errorf("got %s %q", got[0].State, got[0].Detail)
	}
	if uv.gets != 0 {
		t.Error("it began downloading anyway")
	}
}

func TestAFailedPythonInstallLeavesNothingThatLooksInstalled(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{"gcc": true}}, failPip: true}
	in, _ := pythonInstaller(t, gpuNode, sys)
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "Hash mismatch") {
		t.Fatalf("got %s %q", got[0].State, got[0].Detail)
	}
	entries, _ := os.ReadDir(filepath.Join(in.Root, "vllm"))
	for _, e := range entries {
		t.Errorf("left behind: %s", e.Name())
	}
}

func TestAVersionPyPIDoesNotHaveIsSaidSo(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{"gcc": true}}}
	in, _ := pythonInstaller(t, gpuNode, sys)
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "9.9.9"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "PyPI has no vllm 9.9.9") {
		t.Errorf("got %s %q", got[0].State, got[0].Detail)
	}
}