llmgw/internal/mirror/mirror_test.go

package mirror

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func serving(t *testing.T, body string) string {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(body))
	}))
	t.Cleanup(srv.Close)
	return srv.URL + "/model.gguf"
}

func TestAddPlacesAFileInTheVersionedLayout(t *testing.T) {
	root := t.TempDir()
	got, err := Add(context.Background(), root, serving(t, "weights"),
		Entry{Model: "qwen-14b", Version: "2026-09", File: "q4_k_m.gguf"}, nil)
	if err != nil {
		t.Fatal(err)
	}
	if got.Entry.Bytes != 7 {
		t.Errorf("size = %d", got.Entry.Bytes)
	}
	if got.SHA256 == "" {
		t.Error("no checksum reported; a mirror is the one place a corrupt file reaches every node")
	}
	if _, err := os.Stat(filepath.Join(root, "qwen-14b", "2026-09", "q4_k_m.gguf")); err != nil {
		t.Errorf("the file is not where the layout says: %v", err)
	}
}

func TestAddRefusesToReplaceAVersion(t *testing.T) {
	// The whole discipline. A version that can be overwritten takes away the
	// thing versioned paths are for: changing the catalog back and still
	// having the old weights.
	root := t.TempDir()
	e := Entry{Model: "m", Version: "v1", File: "f.gguf"}
	if _, err := Add(context.Background(), root, serving(t, "first"), e, nil); err != nil {
		t.Fatal(err)
	}
	_, err := Add(context.Background(), root, serving(t, "second"), e, nil)
	if !errors.Is(err, ErrExists) {
		t.Fatalf("overwriting a version returned %v", err)
	}
	b, _ := os.ReadFile(e.Path(root))
	if string(b) != "first" {
		t.Errorf("the original was replaced anyway: %q", b)
	}
}

func TestATruncatedDownloadIsNotRenamedIntoPlace(t *testing.T) {
	// A partial file here is a partial file on every node.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Length", "100")
		_, _ = w.Write([]byte("short"))
	}))
	defer srv.Close()

	root := t.TempDir()
	e := Entry{Model: "m", Version: "v1", File: "f.gguf"}
	if _, err := Add(context.Background(), root, srv.URL, e, nil); err == nil {
		t.Fatal("a truncated download was accepted")
	}
	if _, err := os.Stat(e.Path(root)); !os.IsNotExist(err) {
		t.Error("the truncated file was left where a catalog could point at it")
	}
	// And no .part left looking like a model.
	got, _ := List(root)
	if len(got) != 0 {
		t.Errorf("list shows %+v after a failed download", got)
	}
}

func TestPathElementsAreNamesNotPaths(t *testing.T) {
	// A slash would put the file somewhere the layout does not describe, and
	// List would never find it again.
	root := t.TempDir()
	_, err := Add(context.Background(), root, serving(t, "x"),
		Entry{Model: "../escape", Version: "v", File: "f"}, nil)
	if err == nil {
		t.Error("a path was accepted as a name")
	}
}

func TestListIgnoresAnythingNotInTheLayout(t *testing.T) {
	root := t.TempDir()
	if err := os.MkdirAll(filepath.Join(root, "m", "v1"), 0o755); err != nil {
		t.Fatal(err)
	}
	for _, f := range []string{"real.gguf", "half.gguf.part"} {
		if err := os.WriteFile(filepath.Join(root, "m", "v1", f), []byte("x"), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	// A file somebody dropped at the top level, which is not a model.
	if err := os.WriteFile(filepath.Join(root, "notes.txt"), []byte("x"), 0o644); err != nil {
		t.Fatal(err)
	}

	got, err := List(root)
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 1 || got[0].File != "real.gguf" {
		t.Errorf("list = %+v", got)
	}
}

func TestUnusedComparesByPathNotBySpelling(t *testing.T) {
	// A catalog written by hand and a URL built from a path differ in the ways
	// people differ. Reporting a file as unused because somebody typed the
	// base differently would be followed by somebody deleting it.
	entries := []Entry{
		{Model: "m", Version: "v1", File: "a.gguf"},
		{Model: "m", Version: "v2", File: "a.gguf"},
	}
	wanted := []string{
		// A different scheme and a trailing slash on the base, both of which
		// name the same file.
		"https://models.internal/m/v1/a.gguf",
	}
	got := Unused(entries, wanted, "http://models.internal/")
	if len(got) != 1 || got[0].Version != "v2" {
		t.Errorf("unused = %+v, want only v2", got)
	}
}

func TestRemoveTakesTheEmptyDirectoriesWithIt(t *testing.T) {
	// An empty directory left behind makes list show a model that is not
	// there, which is the kind of small lie that costs an afternoon.
	root := t.TempDir()
	e := Entry{Model: "m", Version: "v1", File: "f.gguf"}
	if _, err := Add(context.Background(), root, serving(t, "x"), e, nil); err != nil {
		t.Fatal(err)
	}
	if err := Remove(root, e); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(filepath.Join(root, "m")); !os.IsNotExist(err) {
		t.Error("an empty model directory was left behind")
	}
}

func TestProgressIsReportedForALongDownload(t *testing.T) {
	root := t.TempDir()
	var calls int
	_, err := Add(context.Background(), root, serving(t, strings.Repeat("x", 1<<20)),
		Entry{Model: "m", Version: "v", File: "f"}, func(done, total int64) { calls++ })
	if err != nil {
		t.Fatal(err)
	}
	if calls == 0 {
		t.Error("a download reported no progress at all")
	}
}

func TestANameThatLeavesTheMirrorIsRefused(t *testing.T) {
	root := t.TempDir()
	for _, e := range []Entry{
		{Model: "..", Version: "v", File: "f.gguf"},
		{Model: "m", Version: "..", File: "f.gguf"},
		{Model: "m", Version: "v", File: ".f.sha256"},
		{Model: "a/b", Version: "v", File: "f.gguf"},
	} {
		if _, err := Add(t.Context(), root, "http://127.0.0.1:1/x", e, nil); err == nil ||
			!strings.Contains(err.Error(), "not a usable name") {
			t.Errorf("%+v: %v", e, err)
		}
	}
}

func TestAVersionIsGivenOutOnlyWhole(t *testing.T) {
	root := t.TempDir()
	dir := filepath.Join(root, "m", "v1")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	write := func(name, body string) {
		t.Helper()
		if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	write("config.json", "{}")
	write("model.safetensors.index.json",
		`{"weight_map":{"a":"model-00001-of-00002.safetensors","b":"model-00002-of-00002.safetensors"}}`)
	write("model-00001-of-00002.safetensors", "aa")
	write("model-00002-of-00002.safetensors.part", "b")

	if _, err := ManifestFor(root, "m", "v1"); !errors.Is(err, ErrIncomplete) || !strings.Contains(err.Error(), "still downloading") {
		t.Fatalf("a version with a shard downloading was given out: %v", err)
	}
	// The download failed and its .part went: the index still names it.
	_ = os.Remove(filepath.Join(dir, "model-00002-of-00002.safetensors.part"))
	if _, err := ManifestFor(root, "m", "v1"); !errors.Is(err, ErrIncomplete) ||
		!strings.Contains(err.Error(), "missing model-00002-of-00002.safetensors") {
		t.Fatalf("a version missing a shard its index names was given out: %v", err)
	}
	write("model-00002-of-00002.safetensors", "bb")
	m, err := ManifestFor(root, "m", "v1")
	if err != nil || len(m.Files) != 4 {
		t.Fatalf("the whole version: %+v, %v", m, err)
	}

	// A split GGUF with no index is held to its own numbering.
	gg := filepath.Join(root, "g", "v1")
	_ = os.MkdirAll(gg, 0o755)
	_ = os.WriteFile(filepath.Join(gg, "q-00001-of-00003.gguf"), []byte("x"), 0o644)
	_ = os.WriteFile(filepath.Join(gg, "q-00003-of-00003.gguf"), []byte("x"), 0o644)
	if _, err := ManifestFor(root, "g", "v1"); !errors.Is(err, ErrIncomplete) || !strings.Contains(err.Error(), "q-00002-of-00003.gguf") {
		t.Fatalf("a split missing its second part was given out: %v", err)
	}
}

func TestAMirroredFilesChecksumIsHandedOnAndHiddenFromListings(t *testing.T) {
	root := t.TempDir()
	e := Entry{Model: "m", Version: "v1", File: "f.gguf"}
	added, err := Add(context.Background(), root, serving(t, "weights"), e, nil)
	if err != nil {
		t.Fatal(err)
	}
	entries, _ := List(root)
	if len(entries) != 1 {
		t.Errorf("the checksum was listed as a file of the model: %+v", entries)
	}
	m, err := ManifestFor(root, "m", "v1")
	if err != nil || m.Files[0].SHA256 != added.SHA256 || added.SHA256 == "" {
		t.Errorf("manifest %+v, checksum %q: %v", m, added.SHA256, err)
	}
}