llmgw/internal/agent/provision/strata_test.go

package provision

import (
	"crypto/sha256"
	"encoding/hex"
	"os"
	"path/filepath"
	"strings"
	"testing"

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

// strataRelease is Strata's source archive as GitHub makes one, and its
// release's engine for CUDA 13, with the recipe pointed at the archive.
func strataRelease(t *testing.T, rel *release) {
	t.Helper()
	archive := tarGz(t, map[string]string{
		"Strata-0.2.7/setup.py":         "print('setup')\n",
		"Strata-0.2.7/serve/server.py":  "print('server')\n",
		"Strata-0.2.7/requirements.txt": "numpy==2.5.3; python_version >= \"3.12\"\njinja2==3.1.6\n",
	}, nil)
	rel.assets["v0.2.7.tar.gz"] = archive
	sum := sha256.Sum256(archive)
	engine := tarGz(t, map[string]string{"strata": "\x7fELF", "libcudart.so.13": "lib"}, nil)
	name := "strata-engine-0.1.52-linux-x86_64-cuda-cuda13.tar.gz"
	rel.assets[name] = engine
	esum := sha256.Sum256(engine)
	rel.digests[name] = "sha256:" + hex.EncodeToString(esum[:])

	old := strataReleases
	strataReleases = map[string]struct {
		sha256    string
		published string
		engine    string
	}{"0.2.7": {hex.EncodeToString(sum[:]), "2026-10-27T14:54:42Z", "0.1.52"}}
	t.Cleanup(func() { strataReleases = old })
}

func TestStrataIsInstalledWithItsEngineBesideItsCode(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	in, rel := pythonInstaller(t, turingNode, sys)
	strataRelease(t, rel)

	got := in.Converge(t.Context(), []Want{{Engine: "strata", Version: "0.2.7"}})
	if got[0].State != Ready {
		t.Fatalf("not ready: %+v", got[0])
	}
	dir := filepath.Join(in.Root, "strata", "0.2.7-cuda13")
	if got[0].Binary != filepath.Join(dir, "bin", "python") {
		t.Errorf("binary %s, want the environment's interpreter", got[0].Binary)
	}
	for _, f := range []string{"strata/setup.py", "strata/serve/server.py", "strata-engine/strata", "strata-engine/libcudart.so.13"} {
		if _, err := os.Stat(filepath.Join(dir, f)); err != nil {
			t.Errorf("%s is not in the environment: %v", f, err)
		}
	}
	all := strings.Join(sys.commands, "\n")
	for _, want := range []string{"--exclude-newer 2026-10-28T14:54:42Z", "jinja2==3.1.6"} {
		if !strings.Contains(all, want) {
			t.Errorf("the installation never used %q:\n%s", want, all)
		}
	}
}

func TestStrataRefusesAnEngineWhoseDigestDoesNotMatch(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	in, rel := pythonInstaller(t, turingNode, sys)
	strataRelease(t, rel)
	rel.digests["strata-engine-0.1.52-linux-x86_64-cuda-cuda13.tar.gz"] = "sha256:" + strings.Repeat("0", 64)

	got := in.Converge(t.Context(), []Want{{Engine: "strata", Version: "0.2.7", Variant: "cuda13"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "does not match the checksum") {
		t.Fatalf("an engine that does not match its digest was not refused: %+v", got[0])
	}
}

func TestStrataPicksTheEngineForTheCardsAndDriver(t *testing.T) {
	for _, tc := range []struct {
		name string
		host discovery.Host
		want []string
		err  string
	}{
		{"turing on 580", turingNode, []string{"cuda13", "cuda12"}, ""},
		{"volta", discovery.Host{Arch: "amd64", NVIDIADriver: "580.82.07", MinComputeCapability: 7.0}, []string{"cuda12"}, ""},
		{"driver 575", discovery.Host{Arch: "amd64", NVIDIADriver: "575.51.03", MinComputeCapability: 8.6}, []string{"cuda12"}, ""},
		{"driver 570", discovery.Host{Arch: "amd64", NVIDIADriver: "570.86.10", MinComputeCapability: 8.6}, nil, "575"},
		{"no gpu", discovery.Host{Arch: "amd64"}, nil, "no"},
		{"arm", discovery.Host{Arch: "arm64", NVIDIADriver: "580.82.07", MinComputeCapability: 9.0}, nil, "x86-64"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			bs, err := strataCandidates("0.2.7", "", tc.host)
			if tc.err != "" {
				if err == nil || !strings.Contains(err.Error(), tc.err) {
					t.Fatalf("got %v, %v; want an error saying %q", bs, err, tc.err)
				}
				return
			}
			if err != nil {
				t.Fatal(err)
			}
			var got []string
			for _, b := range bs {
				got = append(got, b.Variant)
			}
			if strings.Join(got, ",") != strings.Join(tc.want, ",") {
				t.Errorf("builds %v, want %v", got, tc.want)
			}
		})
	}
}

func TestStrataIsNotInAnOfflineBundle(t *testing.T) {
	if _, err := Builds("strata", "0.2.7", "amd64"); err == nil {
		t.Fatal("an offline bundle would carry Strata, which needs PyPI")
	}
}