llmgw/internal/hf/hf_test.go

package hf

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"slices"
	"strings"
	"testing"
	"time"
)

// registry stands in for Hugging Face, answering the two endpoints used.
func registry(t *testing.T, h http.HandlerFunc) *Client {
	t.Helper()
	srv := httptest.NewServer(h)
	t.Cleanup(srv.Close)
	return New(srv.URL, time.Second)
}

func TestTheRealSizeOfAWeightsFileIsTheLFSOne(t *testing.T) {
	c := registry(t, func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/tree/abc123") {
			_, _ = w.Write([]byte(`[
			  {"type":"file","path":"README.md","size":900,"oid":"gitblob"},
			  {"type":"file","path":"model.gguf","size":133,
			   "lfs":{"oid":"deadbeef","size":1117320736}}]`))
			return
		}
		_, _ = w.Write([]byte(`{"id":"Qwen/X","sha":"abc123"}`))
	})

	repo, err := c.Repo(t.Context(), "Qwen/X")
	if err != nil {
		t.Fatal(err)
	}
	var gguf File
	for _, f := range repo.Files {
		if f.Path == "model.gguf" {
			gguf = f
		}
	}
	// The plain size of an LFS file is the size of the pointer. Reading that
	// one reports a gigabyte model as 133 bytes, which is the number somebody
	// would then plan disk against.
	if gguf.Bytes != 1117320736 {
		t.Errorf("expected the LFS size, got %d", gguf.Bytes)
	}
	if gguf.SHA256 != "deadbeef" {
		t.Errorf("expected the published checksum, got %q", gguf.SHA256)
	}
}

func TestASmallFileHasNoChecksum(t *testing.T) {
	c := registry(t, func(w http.ResponseWriter, r *http.Request) {
		if strings.Contains(r.URL.Path, "/tree/") {
			_, _ = w.Write([]byte(`[{"type":"file","path":"README.md","size":900,"oid":"gitblob"}]`))
			return
		}
		_, _ = w.Write([]byte(`{"id":"Qwen/X","sha":"abc123"}`))
	})

	repo, _ := c.Repo(t.Context(), "Qwen/X")
	// A git blob hash is not a checksum of the content, so reporting it as one
	// would invite somebody to compare it with what the mirror computes and
	// conclude the download was corrupt.
	if repo.Files[0].SHA256 != "" {
		t.Errorf("a git-stored file should publish no checksum, got %q", repo.Files[0].SHA256)
	}
}

func TestTheDownloadURLIsPinnedToTheRevisionThatWasListed(t *testing.T) {
	r := Repo{ID: "Qwen/X", Revision: "abc123def456"}
	got := r.DownloadURL("https://hf.example", "model.gguf")

	// "main" would mean the mirror fetches whatever is there when it gets round
	// to it, which is not necessarily what was on screen when the button was
	// pressed.
	if !strings.Contains(got, "/resolve/abc123def456/") {
		t.Errorf("the URL should pin the revision; got %s", got)
	}
}

func TestAGatedRepositorySaysSoRatherThanBlamingTheNetwork(t *testing.T) {
	c := registry(t, func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
	})

	_, err := c.Repo(t.Context(), "meta/gated")
	if err == nil || !strings.Contains(err.Error(), "gated or private") {
		t.Errorf("expected the reason to be named; got %v", err)
	}
}

func TestOneGatedResultDoesNotHideThePage(t *testing.T) {
	c := registry(t, func(w http.ResponseWriter, _ *http.Request) {
		// The field is documented as a boolean and is a word in practice.
		// Decoding it strictly fails the whole response, so one gated model in
		// a page of twenty would hide the other nineteen.
		_, _ = w.Write([]byte(`[{"id":"a","gated":false},{"id":"b","gated":"manual"}]`))
	})

	got, err := c.Search(t.Context(), "qwen", 20)
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 2 {
		t.Fatalf("expected both results, got %d", len(got))
	}
	if got[0].Gated.Is() || !got[1].Gated.Is() {
		t.Errorf("gating did not decode: %+v", got)
	}
}

func TestSomethingThatIsNotARepositoryIsRefusedBeforeTheCall(t *testing.T) {
	called := false
	c := registry(t, func(http.ResponseWriter, *http.Request) { called = true })

	if _, err := c.Repo(t.Context(), "qwen"); err == nil {
		t.Error("a bare name is not owner/name and should be refused")
	}
	if called {
		t.Error("nothing should have been requested")
	}
}

func TestWeightsAreToldApartFromDocumentation(t *testing.T) {
	// Crude and right: the registry does not label them, and an operator
	// filling a mirror wants the four files out of thirty worth gigabytes.
	for _, f := range []string{"model.gguf", "model-00001-of-00002.safetensors", "pytorch_model.bin"} {
		if !(File{Path: f}).IsWeights() {
			t.Errorf("%s should count as weights", f)
		}
	}
	for _, f := range []string{"README.md", "config.json", ".gitattributes", "tokenizer.model"} {
		if (File{Path: f}).IsWeights() {
			t.Errorf("%s should not count as weights", f)
		}
	}
}

func TestShardsBecomeOneSetAndAlternativesStaySeparate(t *testing.T) {
	// Straight from a real repository, which holds both the split and the
	// joined copy of the same quantisation.
	sets := Group([]File{
		{Path: "qwen-fp16-00001-of-00004.gguf", Bytes: 1},
		{Path: "qwen-fp16-00002-of-00004.gguf", Bytes: 1},
		{Path: "qwen-fp16-00003-of-00004.gguf", Bytes: 1},
		{Path: "qwen-fp16-00004-of-00004.gguf", Bytes: 1},
		{Path: "qwen-fp16.gguf", Bytes: 4},
		{Path: "qwen-q4_k_m.gguf", Bytes: 2},
	})

	if len(sets) != 3 {
		t.Fatalf("expected the shards as one and the two whole files apart, got %d: %+v",
			len(sets), sets)
	}
	if sets[0].Shards != 4 || sets[0].Bytes != 4 {
		t.Errorf("the shard set should total its parts, got %+v", sets[0])
	}
	// Two GGUFs in a repository are alternatives: fetching both is wasting
	// gigabytes on a copy nobody asked for.
	if sets[1].Shards != 0 || sets[2].Shards != 0 {
		t.Errorf("whole files are not shard sets: %+v %+v", sets[1], sets[2])
	}
}

func TestASplitAndAJoinedCopyDoNotReadAsOneEntry(t *testing.T) {
	sets := Group([]File{
		{Path: "m-q4-00001-of-00002.gguf"}, {Path: "m-q4-00002-of-00002.gguf"},
		{Path: "m-q4.gguf"},
	})
	if sets[0].Name == sets[1].Name {
		t.Errorf("the split and the joined copy collided on %q", sets[0].Name)
	}
	if !strings.Contains(sets[0].Name, "-of-00002") {
		t.Errorf("a shard set should keep the notation; got %q", sets[0].Name)
	}
}

func TestAShardSetPublishesNoSingleChecksum(t *testing.T) {
	sets := Group([]File{
		{Path: "m-00001-of-00002.safetensors", SHA256: "aaa"},
		{Path: "m-00002-of-00002.safetensors", SHA256: "bbb"},
	})
	// Several files have several checksums and no single value for the set;
	// showing one invites comparing it against the wrong thing.
	if got := sets[0].SHA256(); got != "" {
		t.Errorf("expected no checksum for a set, got %q", got)
	}
}

func TestExploringAsksForWhatThePageFilters(t *testing.T) {
	var got url.Values
	calls := 0
	c := registry(t, func(w http.ResponseWriter, r *http.Request) {
		calls++
		got = r.URL.Query()
		_, _ = w.Write([]byte(`[{"id":"a/b-GGUF","gguf":{"total":7615616512,"architecture":"qwen2","context_length":32768}}]`))
	})

	q := Query{Sort: SortTrending, Task: "text-generation", GGUF: true, MaxParams: 9e9, Limit: 30}
	models, err := c.Explore(t.Context(), q)
	if err != nil {
		t.Fatal(err)
	}
	for k, want := range map[string]string{
		"sort": "trendingScore", "pipeline_tag": "text-generation", "filter": "gguf",
		"num_parameters": "min:0,max:9000000000", "limit": "30",
	} {
		if got.Get(k) != want {
			t.Errorf("%s = %q, want %q", k, got.Get(k), want)
		}
	}
	// Expanding any field returns only the expanded ones, so the listing has
	// to name the counts it shows as well as the GGUF summary.
	if !slices.Contains(got["expand[]"], "gguf") || !slices.Contains(got["expand[]"], "downloads") {
		t.Errorf("expand = %v", got["expand[]"])
	}
	if len(models) != 1 || models[0].Params() != 7615616512 || models[0].Architecture() != "qwen2" {
		t.Errorf("the GGUF summary did not decode: %+v", models)
	}

	// The page opens on this; a second look within minutes is not a second
	// trip to the registry.
	if _, err := c.Explore(t.Context(), q); err != nil {
		t.Fatal(err)
	}
	if calls != 1 {
		t.Errorf("the registry was asked %d times", calls)
	}
}

func TestARepositoryCarriesItsCardWithoutTheHeader(t *testing.T) {
	c := registry(t, func(w http.ResponseWriter, r *http.Request) {
		switch {
		case strings.HasPrefix(r.URL.Path, "/api/models/q/m/tree/"):
			_, _ = w.Write([]byte(`[{"type":"file","path":"README.md","size":40},{"type":"file","path":"m.gguf","size":1}]`))
		case r.URL.Path == "/api/models/q/m":
			_, _ = w.Write([]byte(`{"id":"q/m","sha":"abc","cardData":{"license":"apache-2.0","base_model":"Qwen/Qwen2.5-0.5B","language":["en","tr"]}}`))
		case r.URL.Path == "/q/m/raw/abc/README.md":
			_, _ = w.Write([]byte("---\nlicense: apache-2.0\n---\n\n# M\n\nA model.\n"))
		default:
			http.NotFound(w, r)
		}
	})
	got, err := c.Repo(t.Context(), "q/m")
	if err != nil {
		t.Fatal(err)
	}
	if got.Card != "# M\n\nA model." {
		t.Errorf("card = %q", got.Card)
	}
	if got.License != "apache-2.0" || len(got.BaseModel) != 1 || got.BaseModel[0] != "Qwen/Qwen2.5-0.5B" ||
		len(got.Languages) != 2 {
		t.Errorf("card header = %+v", got)
	}
}

func TestAListingNamesTheBaseModelAndLicenseFromItsTags(t *testing.T) {
	m := Model{Tags: []string{"gguf", "base_model:quantized:Qwen/Qwen3-8B", "base_model:Qwen/Qwen3-8B", "license:apache-2.0"}}
	if m.BaseModel() != "Qwen/Qwen3-8B" || m.License() != "apache-2.0" {
		t.Errorf("base %q, license %q", m.BaseModel(), m.License())
	}
	if (Model{Tags: []string{"base_model:org/x"}}).BaseModel() != "org/x" {
		t.Error("a base model with no relation should read as itself")
	}
}