llmgw/internal/agent/discovery/weights_test.go

package discovery

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

func write(t *testing.T, path string, size int) {
	t.Helper()
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(path, make([]byte, size), 0o644); err != nil {
		t.Fatal(err)
	}
}

func TestScanFindsFilesAndDirectoriesOfWeights(t *testing.T) {
	dir := t.TempDir()
	write(t, filepath.Join(dir, "qwen-3b-q4.gguf"), 1024)
	// A HuggingFace checkout: a directory of shards rather than one file.
	write(t, filepath.Join(dir, "Llama-3.3-70B", "model-00001.safetensors"), 2048)
	write(t, filepath.Join(dir, "Llama-3.3-70B", "config.json"), 64)

	got := ScanWeights([]string{dir})
	if len(got) != 2 {
		t.Fatalf("found %d entries: %+v", len(got), got)
	}
	byName := map[string]int64{}
	for _, w := range got {
		byName[filepath.Base(w.Path)] = w.SizeBytes
	}
	if byName["qwen-3b-q4.gguf"] != 1024 {
		t.Errorf("gguf size = %d", byName["qwen-3b-q4.gguf"])
	}
	// The directory's size is everything in it, config included: what matters
	// is how much would have to be moved to put it somewhere else.
	if byName["Llama-3.3-70B"] != 2048+64 {
		t.Errorf("checkout size = %d, want %d", byName["Llama-3.3-70B"], 2048+64)
	}
}

func TestAnEmptyDirectoryIsNotWeights(t *testing.T) {
	// A half-finished download looks exactly like this, and reporting it as
	// present would have placement choose a node whose engine then fails to
	// load — the failure this whole feature exists to avoid.
	dir := t.TempDir()
	if err := os.MkdirAll(filepath.Join(dir, "downloading"), 0o755); err != nil {
		t.Fatal(err)
	}
	write(t, filepath.Join(dir, "downloading", "README.md"), 10)

	if got := ScanWeights([]string{dir}); len(got) != 0 {
		t.Errorf("a directory with no weight files was reported: %+v", got)
	}
}

func TestNotesAndScriptsAreNotReportedAsModels(t *testing.T) {
	dir := t.TempDir()
	write(t, filepath.Join(dir, "notes.txt"), 10)
	write(t, filepath.Join(dir, "fetch.sh"), 10)
	write(t, filepath.Join(dir, "real.gguf"), 10)

	got := ScanWeights([]string{dir})
	if len(got) != 1 || filepath.Base(got[0].Path) != "real.gguf" {
		t.Errorf("scan returned %+v", got)
	}
}

func TestAMissingDirectoryIsNotAnError(t *testing.T) {
	// A mount that has not come back yet looks exactly like a directory that is
	// not there, and the node is simply one without those weights until it is.
	got := ScanWeights([]string{filepath.Join(t.TempDir(), "never-existed")})
	if len(got) != 0 {
		t.Errorf("got %+v", got)
	}
}

func TestScanningNothingIsTheDefaultAndConstrainsNothing(t *testing.T) {
	// What makes the feature optional: an agent with no directories configured
	// behaves exactly as every agent did before this existed.
	if got := ScanWeights(nil); len(got) != 0 {
		t.Errorf("got %+v", got)
	}
	if got := ScanWeights([]string{""}); len(got) != 0 {
		t.Errorf("an empty directory name produced %+v", got)
	}
}

func TestTheSamePathIsReportedOnce(t *testing.T) {
	// Two configured directories can overlap, or be the same mount reached two
	// ways. Counting a model twice would overstate what a fabric holds.
	dir := t.TempDir()
	write(t, filepath.Join(dir, "m.gguf"), 10)
	if got := ScanWeights([]string{dir, dir}); len(got) != 1 {
		t.Errorf("the same directory twice produced %+v", got)
	}
}

func TestLocalStorageIsNeverMistakenForShared(t *testing.T) {
	// The conservative direction. A shared filer read as local costs an
	// unnecessary copy; local read as shared would have the cluster believe a
	// model is everywhere when it is on one node.
	dir := t.TempDir()
	write(t, filepath.Join(dir, "m.gguf"), 10)
	got := ScanWeights([]string{dir})
	if len(got) != 1 {
		t.Fatalf("got %+v", got)
	}
	if got[0].Shared {
		t.Error("a temporary directory on local disk was reported as shared")
	}
}