llmgw/internal/agent/weights/set_test.go

package weights

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// mirrorOf serves a manifest at "/m/v/" and the files beneath it.
func mirrorOf(t *testing.T, files map[string]string) (*Fetcher, *httptest.Server, *int) {
	t.Helper()
	hits := 0
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/") {
			var b strings.Builder
			b.WriteString(`{"model":"m","version":"v","files":[`)
			first := true
			total := 0
			for name, body := range files {
				if !first {
					b.WriteString(",")
				}
				first = false
				fmt.Fprintf(&b, `{"name":%q,"bytes":%d}`, name, len(body))
				total += len(body)
			}
			fmt.Fprintf(&b, `],"bytes":%d}`, total)
			_, _ = w.Write([]byte(b.String()))
			return
		}
		body, ok := files[strings.TrimPrefix(r.URL.Path, "/m/v/")]
		if !ok {
			http.NotFound(w, r)
			return
		}
		hits++
		_, _ = w.Write([]byte(body))
	}))
	t.Cleanup(srv.Close)
	return &Fetcher{CacheDir: t.TempDir()}, srv, &hits
}

func TestASetIsFetchedIntoOneDirectory(t *testing.T) {
	f, srv, _ := mirrorOf(t, map[string]string{
		"model-00001-of-00002.safetensors": "aaaa",
		"model-00002-of-00002.safetensors": "bb",
		"config.json":                      "{}",
	})

	got, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
	if err != nil {
		t.Fatal(err)
	}
	st, err := os.Stat(got)
	if err != nil || !st.IsDir() {
		t.Fatalf("expected a directory, got %q: %v", got, err)
	}
	// A directory, because that is what a multi-file model is: vLLM reads the
	// shards and the config out of the folder, and a catalog entry naming one
	// file could express neither.
	for _, name := range []string{"model-00001-of-00002.safetensors", "config.json"} {
		if _, err := os.Stat(filepath.Join(got, name)); err != nil {
			t.Errorf("%s did not arrive: %v", name, err)
		}
	}
}

func TestAFinishedSetIsNotFetchedAgain(t *testing.T) {
	f, srv, hits := mirrorOf(t, map[string]string{"a.gguf": "aaaa", "b.gguf": "bb"})

	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err != nil {
		t.Fatal(err)
	}
	after := *hits
	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err != nil {
		t.Fatal(err)
	}
	if *hits != after {
		t.Errorf("the second call downloaded again: %d then %d", after, *hits)
	}
}

func TestAHalfFetchedSetIsNotTreatedAsFinished(t *testing.T) {
	f, srv, _ := mirrorOf(t, map[string]string{"a.gguf": "aaaa", "b.gguf": "bb"})
	dir, err := f.cachePath(srv.URL + "/m/v")
	if err != nil {
		t.Fatal(err)
	}
	// One file present, no completion marker — an interrupted download.
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "a.gguf"), []byte("aaaa"), 0o644); err != nil {
		t.Fatal(err)
	}

	got, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
	if err != nil {
		t.Fatal(err)
	}
	// A directory with three of four shards loads as a broken model rather
	// than failing to load, so the marker is what says the set is whole.
	if _, err := os.Stat(filepath.Join(got, "b.gguf")); err != nil {
		t.Errorf("the missing file was not fetched: %v", err)
	}
}

func TestAFileAlreadyThereAtTheRightSizeIsNotRefetched(t *testing.T) {
	f, srv, hits := mirrorOf(t, map[string]string{"a.gguf": "aaaa", "b.gguf": "bb"})
	dir, _ := f.cachePath(srv.URL + "/m/v")
	_ = os.MkdirAll(dir, 0o755)
	_ = os.WriteFile(filepath.Join(dir, "a.gguf"), []byte("aaaa"), 0o644)

	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err != nil {
		t.Fatal(err)
	}
	// Resuming a set means not paying for the shards that already landed.
	if *hits != 1 {
		t.Errorf("expected only the missing file to be fetched, got %d downloads", *hits)
	}
}

func TestAFilenameThatEscapesTheDirectoryIsRefused(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(`{"files":[{"name":"../../etc/passwd","bytes":1}]}`))
	}))
	defer srv.Close()
	f := &Fetcher{CacheDir: t.TempDir()}

	_, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
	// The mirror is a place an operator copies files into by hand, so what it
	// reports is not automatically safe to join onto a path.
	if err == nil || !strings.Contains(err.Error(), "offered a file named") {
		t.Errorf("expected the name to be refused, got %v", err)
	}
}

func TestADirectoryListingIsNotAManifest(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte("<pre><a href=\"a.gguf\">a.gguf</a></pre>"))
	}))
	defer srv.Close()
	f := &Fetcher{CacheDir: t.TempDir()}

	_, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
	// "invalid character '<'" sends somebody looking at the wrong thing.
	if err == nil || !strings.Contains(err.Error(), "has to be a mirror version") {
		t.Errorf("the reason should name the cause, got %v", err)
	}
}

func TestOnlyATrailingSlashMeansASet(t *testing.T) {
	for src, want := range map[string]bool{
		"http://mirror/m/v/":           true,
		"http://mirror/m/v/model.gguf": false,
		"/var/lib/models/":             false, // a local directory is not fetched
		"Qwen/Qwen2.5-7B-Instruct":     false, // a repo id, resolved by the engine
	} {
		if got := IsSet(src); got != want {
			t.Errorf("IsSet(%q) = %v, want %v", src, got, want)
		}
	}
}

func TestASetFileThatDoesNotMatchTheMirrorsChecksumIsRefused(t *testing.T) {
	good := "the shard as mirrored"
	sum := sha256.Sum256([]byte(good))
	served := "the shard damaged on the way"
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/") {
			fmt.Fprintf(w, `{"model":"m","version":"v","files":[{"name":"s.safetensors","bytes":%d,"sha256":%q}]}`,
				len(served), hex.EncodeToString(sum[:]))
			return
		}
		_, _ = w.Write([]byte(served))
	}))
	t.Cleanup(srv.Close)
	f := &Fetcher{CacheDir: t.TempDir()}
	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err == nil || !strings.Contains(err.Error(), "sha256") {
		t.Fatalf("a damaged shard was accepted: %v", err)
	}
	served = good
	// The size differs, so it is fetched again; this time it matches.
	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/"); err != nil {
		t.Fatalf("the shard as mirrored: %v", err)
	}
}

func TestASetStillDownloadingOnTheMirrorSaysSo(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		http.Error(w, "m/v: this version is not complete: still downloading model-00002-of-00002.safetensors", http.StatusConflict)
	}))
	t.Cleanup(srv.Close)
	f := &Fetcher{CacheDir: t.TempDir()}
	_, err := f.Ensure(t.Context(), srv.URL+"/m/v/")
	if err == nil || !strings.Contains(err.Error(), "not complete on the mirror yet") ||
		!strings.Contains(err.Error(), "model-00002-of-00002") {
		t.Fatalf("got %v", err)
	}
}