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")
	}
}