llmgw/internal/cli/mirror_test.go

package cli

import (
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
)

func TestTheMirrorRefusesEverythingButReads(t *testing.T) {
	// Stated here rather than inherited from http.FileServer, so a handler
	// added later cannot quietly make the mirror writable — which would turn
	// it into a way to replace a model with another one that every node then
	// fetches without complaint.
	h := readOnly(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))

	for _, m := range []string{http.MethodGet, http.MethodHead} {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest(m, "/x.gguf", nil))
		if rec.Code != http.StatusOK {
			t.Errorf("%s was refused: %d", m, rec.Code)
		}
	}
	for _, m := range []string{http.MethodPut, http.MethodPost, http.MethodDelete, http.MethodPatch} {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest(m, "/x.gguf", nil))
		if rec.Code != http.StatusMethodNotAllowed {
			t.Errorf("%s was allowed: %d", m, rec.Code)
		}
	}
}

func TestTheFilenameIsTakenFromTheURLWhenNotGiven(t *testing.T) {
	for _, c := range []struct{ url, want string }{
		{"https://h/a/b/model.gguf", "model.gguf"},
		{"https://h/model.gguf?download=true", "model.gguf"},
		{"model.gguf", "model.gguf"},
	} {
		if got := mirror.LastPathElement(c.url); got != c.want {
			t.Errorf("mirror.LastPathElement(%q) = %q, want %q", c.url, got, c.want)
		}
	}
}

func TestAVersionAnswersHeadAsWellAsGet(t *testing.T) {
	// A sharded model's catalog source is the version directory, and the serve
	// path checks a source with HEAD before accepting it. If the manifest
	// answered only GET, every multi-file model would be refused as missing.
	root := t.TempDir()
	if err := os.MkdirAll(filepath.Join(root, "llama", "v1"), 0o755); err != nil {
		t.Fatal(err)
	}
	for _, name := range []string{"model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"} {
		if err := os.WriteFile(filepath.Join(root, "llama", "v1", name), []byte("x"), 0o644); err != nil {
			t.Fatal(err)
		}
	}

	h := readOnly(withManifest(root, http.FileServer(http.Dir(root))))
	w := httptest.NewRecorder()
	h.ServeHTTP(w, httptest.NewRequest(http.MethodHead, "/llama/v1/", nil))

	if w.Code != http.StatusOK {
		t.Fatalf("a version should answer HEAD; got %d", w.Code)
	}
	if got := w.Header().Get("Content-Type"); got != "application/json" {
		t.Errorf("expected a manifest, got %q", got)
	}
}

func TestTheMirrorDoesNotServeItsBookkeeping(t *testing.T) {
	root := t.TempDir()
	dir := filepath.Join(root, "llama", "v1")
	_ = os.MkdirAll(dir, 0o755)
	_ = os.WriteFile(filepath.Join(dir, "m.gguf"), []byte("x"), 0o644)
	_ = os.WriteFile(filepath.Join(dir, ".m.gguf.source"), []byte("https://example.com/m.gguf\n"), 0o644)

	h := readOnly(withManifest(root, http.FileServer(http.Dir(root))))
	for path, want := range map[string]int{
		"/llama/v1/m.gguf":         http.StatusOK,
		"/llama/v1/.m.gguf.source": http.StatusNotFound,
		"/llama/.hidden/x":         http.StatusNotFound,
	} {
		w := httptest.NewRecorder()
		h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, path, nil))
		if w.Code != want {
			t.Errorf("%s: %d, expected %d", path, w.Code, want)
		}
	}
}