internal/hub/hub_test.go
Rama main · réplica de lectura del repositorio original.
package hub
import (
"path/filepath"
"strings"
"testing"
)
func TestClassify_ShardGroup(t *testing.T) {
cases := []struct {
path string
shard, shards int
group string
}{
{"Qwen3-8B-Q4_K_M.gguf", 0, 0, "Qwen3-8B-Q4_K_M.gguf"},
{"model-IQ3_XXS.gguf", 0, 0, "model-IQ3_XXS.gguf"},
{"Llama-3-70B.BF16.gguf", 0, 0, "Llama-3-70B.BF16.gguf"},
// Every shard of one file collapses onto a single group key, and that key
// is what the picker shows as the row's name — the shared filename with
// the part numbering removed.
{"Q8_0/big-Q8_0-00001-of-00003.gguf", 1, 3, "Q8_0/big-Q8_0.gguf"},
{"Q8_0/big-Q8_0-00003-of-00003.gguf", 3, 3, "Q8_0/big-Q8_0.gguf"},
}
for _, c := range cases {
f := File{Path: c.path}
classify(&f)
if f.Shard != c.shard || f.Shards != c.shards || f.Group != c.group {
t.Errorf("classify(%q) = shard %d/%d group %q; want %d/%d %q",
c.path, f.Shard, f.Shards, f.Group, c.shard, c.shards, c.group)
}
}
}
func TestClassify_Projector(t *testing.T) {
// A projector is a companion file, not a model: the flag is what sorts it
// below the weights and replaces its fit verdict with "companion".
for _, p := range []string{
"mmproj-F16.gguf",
"mmproj-model-f16.gguf",
"Qwen2-VL-7B-Instruct.mmproj.gguf",
"sub/Qwen2-VL-mmproj.gguf",
"MMPROJ-BF16.GGUF",
"llava-v1.6-mm_proj-Q8_0.gguf",
} {
f := File{Path: p}
classify(&f)
if !f.Projector {
t.Errorf("classify(%q).Projector = false, want true", p)
}
}
for _, p := range []string{"Qwen3-8B-Q4_K_M.gguf", "model-IQ3_XXS.gguf", "projection-tuned-Q8_0.gguf"} {
f := File{Path: p}
classify(&f)
if f.Projector {
t.Errorf("classify(%q).Projector = true, want false", p)
}
}
}
func TestIsModelFile(t *testing.T) {
for _, p := range []string{"a.gguf", "sub/B.GGUF", "sam.ggml"} {
if !IsModelFile(p) {
t.Errorf("IsModelFile(%q) = false, want true", p)
}
}
for _, p := range []string{"README.md", "config.json", "model.safetensors", "preview.png"} {
if IsModelFile(p) {
t.Errorf("IsModelFile(%q) = true, want false", p)
}
}
}
// A repo that ships any GGUF stays a quant repo: its configs, tokenizer and
// README are noise beside the files a loader can open.
func TestSelectFiles_QuantRepo(t *testing.T) {
files := []File{
{Path: "Qwen3-8B-Q4_K_M.gguf"},
{Path: "Qwen3-8B-Q5_K_M.gguf"},
{Path: "mmproj-F16.gguf"},
{Path: "config.json"},
{Path: "tokenizer.json"},
{Path: "model.safetensors"},
{Path: "README.md"},
{Path: ".gitattributes"},
}
got := SelectFiles("unsloth/Qwen3-8B-GGUF", files)
want := []string{"Qwen3-8B-Q4_K_M.gguf", "Qwen3-8B-Q5_K_M.gguf", "mmproj-F16.gguf"}
if len(got) != len(want) {
t.Fatalf("got %d files, want %d: %+v", len(got), len(want), got)
}
for i, p := range want {
if got[i].Path != p {
t.Errorf("file[%d] = %q, want %q", i, got[i].Path, p)
}
}
}
// A TRELLIS.2 repo is a component set: the safetensors and the manifests that
// name them ARE the model, they land on disk as one folder, and every one of
// them shares a group key so the picker offers the set as a single row. The
// transformers scaffolding (tokenizer, preprocessor, chat template) stays out.
func TestSelectFiles_ComponentSet(t *testing.T) {
files := []File{
{Path: "pipeline.json"},
{Path: "texturing_pipeline.json"},
{Path: "ckpts/shape_slat_flow.safetensors"},
{Path: "ckpts/shape_slat_flow.json"},
{Path: "ckpts/texture_flow.safetensors"},
{Path: "ckpts/texture_flow.json"},
{Path: "config.json"},
{Path: "tokenizer.json"},
{Path: "preprocessor_config.json"},
{Path: "chat_template.jinja"},
{Path: "README.md"},
{Path: "images/teaser.png"},
}
got := SelectFiles("microsoft/TRELLIS.2-4B", files)
want := map[string]bool{
"pipeline.json": true,
"texturing_pipeline.json": true,
"ckpts/shape_slat_flow.safetensors": true,
"ckpts/shape_slat_flow.json": true,
"ckpts/texture_flow.safetensors": true,
"ckpts/texture_flow.json": true,
"config.json": true,
}
if len(got) != len(want) {
t.Fatalf("got %d files, want %d: %+v", len(got), len(want), got)
}
for _, f := range got {
if !want[f.Path] {
t.Errorf("unexpected file %q", f.Path)
}
if f.Group != "microsoft/TRELLIS.2-4B" {
t.Errorf("%q: group = %q, want the repo id so the set is one picker row", f.Path, f.Group)
}
}
}
// Neither a quant file nor a component file: nothing can be loaded, so the
// picker must be empty rather than offering a README as a download.
func TestSelectFiles_UnloadableRepo(t *testing.T) {
got := SelectFiles("someone/notes", []File{{Path: "README.md"}, {Path: "train.py"}, {Path: "tokenizer.json"}})
if len(got) != 0 {
t.Fatalf("want no files, got %+v", got)
}
}
func TestHF_CheckURL(t *testing.T) {
h := NewHF()
ok := []string{
"https://huggingface.co/o/r/resolve/main/a.gguf",
"https://cdn-lfs-us-1.hf.co/repos/x/y",
"https://cdn-lfs.huggingface.co/repos/x",
}
for _, u := range ok {
if err := h.CheckURL(u); err != nil {
t.Errorf("CheckURL(%q) = %v, want nil", u, err)
}
}
// A redirect or a poisoned API response must not be able to walk the
// download off the hub, and plain http must not be accepted either.
bad := []string{
"http://huggingface.co/o/r/resolve/main/a.gguf",
"https://evil.com/a.gguf",
"https://huggingface.co.evil.com/a.gguf",
"https://nothf.co.uk/a.gguf",
}
for _, u := range bad {
if err := h.CheckURL(u); err == nil {
t.Errorf("CheckURL(%q) = nil, want refusal", u)
}
}
}
func TestValidRepoIDAndPath(t *testing.T) {
for _, id := range []string{"unsloth/Qwen3-8B-GGUF", "a/b.c-d_e"} {
if err := validRepoID(id); err != nil {
t.Errorf("validRepoID(%q) = %v, want nil", id, err)
}
}
for _, id := range []string{"noslash", "a/b/c", "/b", "a/", "../etc", "a/..", "o/r?x=1", "o/r#f"} {
if err := validRepoID(id); err == nil {
t.Errorf("validRepoID(%q) = nil, want refusal", id)
}
}
for _, p := range []string{"a.gguf", "Q4_K_M/a-00001-of-00002.gguf"} {
if err := validRepoPath(p); err != nil {
t.Errorf("validRepoPath(%q) = %v, want nil", p, err)
}
}
for _, p := range []string{"", "/abs.gguf", "../a.gguf", "a/../../b", "a\\b.gguf", "a//b"} {
if err := validRepoPath(p); err == nil {
t.Errorf("validRepoPath(%q) = nil, want refusal", p)
}
}
}
func TestFileURL_Escapes(t *testing.T) {
h := NewHF()
got, err := h.FileURL("unsloth/Qwen3-8B-GGUF", "Q4_K_M/model file.gguf")
if err != nil {
t.Fatal(err)
}
want := "https://huggingface.co/unsloth/Qwen3-8B-GGUF/resolve/main/Q4_K_M/model%20file.gguf"
if got != want {
t.Errorf("FileURL = %q, want %q", got, want)
}
if _, err := h.FileURL("unsloth/Qwen3", "../../etc/passwd"); err == nil {
t.Error("FileURL accepted a traversing path")
}
}
func TestRepoDirName(t *testing.T) {
// One folder per repo, nested under the publisher — the hub's own layout.
cases := map[string]string{
"unsloth/Qwen3-8B-GGUF": filepath.Join("unsloth", "Qwen3-8B-GGUF"),
"a b/c:d": filepath.Join("a-b", "c-d"),
"noslash": "noslash",
// A traversing owner cannot climb out: every segment is sanitized before
// it is joined, so ".." never survives as ".."
"../..//etc": filepath.Join("--", "..--etc"),
}
for in, want := range cases {
if got := RepoDirName(in); got != want {
t.Errorf("RepoDirName(%q) = %q, want %q", in, got, want)
}
}
}
func TestParamsB(t *testing.T) {
cases := map[string]float64{
"unsloth/Qwen3-235B-A22B-GGUF": 235, // MoE: total, never the active count
"bartowski/Meta-Llama-3.1-70B-GGUF": 70,
"unsloth/Qwen3.6-35B-A3B-Q4_K_M": 35,
"TheBloke/Mixtral-8x7B-v0.1-GGUF": 56, // product; overstates, which is the safe way to be wrong
"unsloth/Qwen3-4B-Instruct-GGUF": 4,
"someone/tinyllama-1.1b-chat-gguf": 1.1,
"ggml-org/gpt-oss-120b-GGUF": 120,
"unsloth/Qwen3-8B-GGUF-Q4_K_M": 8, // a quant tag is not a size
"mradermacher/some-model-i1-GGUF": 0, // no size stated -> unknown
"openai/whisper-large-v3": 0,
}
for in, want := range cases {
if got := ParamsB(in); got != want {
t.Errorf("ParamsB(%q) = %v, want %v", in, got, want)
}
}
}
func TestWithinParams(t *testing.T) {
if !WithinParams("unsloth/Qwen3-235B-A22B-GGUF", 0) {
t.Error("a zero cap must not filter anything")
}
if WithinParams("unsloth/Qwen3-235B-A22B-GGUF", 120) {
t.Error("235B passed a 120B cap")
}
if !WithinParams("unsloth/Qwen3-70B-GGUF", 120) {
t.Error("70B failed a 120B cap")
}
if !WithinParams("mradermacher/some-model-GGUF", 120) {
t.Error("an unknown size must be kept, not hidden")
}
if WithinParams("ggml-org/gpt-oss-120b-GGUF", 120) {
t.Error("the cap is exclusive: 120B must not pass a 120B cap")
}
}
// The Video tab is the one category no single HF query describes: the exact
// pipeline tags are split across the catalog and HF ANDs its filters, so it
// asks twice and merges. `gguf` is deliberately absent — the useful video repos
// are LoRAs and safetensors component sets.
func TestSearchFilterSets_Video(t *testing.T) {
got := searchFilterSets("video")
if len(got) != 2 {
t.Fatalf("searchFilterSets(video) = %v, want two sets", got)
}
for _, set := range got {
if len(set) != 1 {
t.Fatalf("set %v: want one filter per set, so neither hides the other", set)
}
if set[0] == "gguf" {
t.Fatal("video must not require gguf: LoRA and safetensors repos are the point")
}
}
if got[0][0] != "text-to-video" || got[1][0] != "image-to-video" {
t.Fatalf("searchFilterSets(video) = %v", got)
}
// Every other category is one set, unchanged.
if sets := searchFilterSets("image"); len(sets) != 1 || len(sets[0]) != 2 {
t.Fatalf("searchFilterSets(image) = %v, want one gguf+pipeline set", sets)
}
}
// Decision models share no pipeline tag, so the tab unions three tag queries,
// each still requiring a GGUF.
func TestSearchFilterSets_Decision(t *testing.T) {
got := searchFilterSets("decision")
if len(got) != 3 {
t.Fatalf("searchFilterSets(decision) = %v, want three sets", got)
}
for i, want := range []string{"decision-model", "typed-decisions", "system-one"} {
if len(got[i]) != 2 || got[i][0] != "gguf" || got[i][1] != want {
t.Errorf("set %d = %v, want [gguf %s]", i, got[i], want)
}
}
}
// A decision repo tagged text-generation is still not an LLM.
func TestIsText_DecisionRepoExcluded(t *testing.T) {
if isText(Model{Pipeline: "text-generation", Tags: []string{"gguf", "decision-model"}}) {
t.Error("decision-model repo kept in the LLM tab")
}
if !isText(Model{Pipeline: "text-generation", Tags: []string{"gguf"}}) {
t.Error("plain text-generation repo dropped")
}
}
// A LoRA repo is a pile of ALTERNATIVES, not the parts of one model: grouping
// it under the repo id offered a single row that downloaded all seventeen.
func TestSelectFiles_LooseWeightsRepo(t *testing.T) {
files := []File{
{Path: "README.md"},
{Path: "minimax_h3_fl2v_turbo_4step_v1.1_768p_bf16.safetensors"},
{Path: "minimax_h3_fl2v_turbo_4step_v1.1_768p_fp8.safetensors"},
{Path: "minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors"},
}
for i := range files {
classify(&files[i])
}
got := SelectFiles("lightx2v/Minimax-h3-Turbo", files)
if len(got) != 3 {
t.Fatalf("kept %d files, want the 3 safetensors: %v", len(got), got)
}
groups := map[string]bool{}
for _, f := range got {
if f.Group == "lightx2v/Minimax-h3-Turbo" {
t.Fatalf("%s: grouped as the whole repo, want its own row", f.Path)
}
if f.Group == "" {
t.Fatalf("%s: no group key", f.Path)
}
groups[f.Group] = true
}
if len(groups) != 3 {
t.Fatalf("3 alternatives collapsed onto %d rows", len(groups))
}
}
// MarkSelection lists the WHOLE repo: the curated files keep SelectFiles'
// grouping, everything else comes back flagged and on its own row.
func TestMarkSelection_ListsEverything(t *testing.T) {
files := []File{
{Path: "README.md"},
{Path: "Qwen3-8B-Q4_K_M.gguf"},
{Path: "mmproj-F16.gguf"},
{Path: "tokenizer.json"},
{Path: "vae/diffusion_pytorch_model.safetensors"},
}
got := MarkSelection("unsloth/Qwen3-8B-GGUF", files)
if len(got) != len(files) {
t.Fatalf("listed %d files, want all %d", len(got), len(files))
}
aux := map[string]bool{}
for _, f := range got {
aux[f.Path] = f.Aux
if f.Aux && f.Group != f.Path {
t.Errorf("%s: aux file grouped as %q, want its own path", f.Path, f.Group)
}
}
for p, want := range map[string]bool{
"Qwen3-8B-Q4_K_M.gguf": false,
"mmproj-F16.gguf": false,
"README.md": true,
"tokenizer.json": true,
"vae/diffusion_pytorch_model.safetensors": true, // a quant repo: the gguf is the model
} {
if aux[p] != want {
t.Errorf("%s: aux=%v, want %v", p, aux[p], want)
}
}
}
// A repo with nothing loadable is still fully listed — that is the case the
// old picker answered with "this repo carries no GGUF files" and no recourse.
func TestMarkSelection_UnloadableRepo(t *testing.T) {
got := MarkSelection("someone/notes", []File{{Path: "README.md"}, {Path: "train.py"}})
if len(got) != 2 {
t.Fatalf("listed %d files, want 2", len(got))
}
for _, f := range got {
if !f.Aux {
t.Errorf("%s: want aux", f.Path)
}
}
}
// Rows as the hub returned them under filter=gguf: the LLM tab must drop the
// ones that say they are another job, and keep the untagged quant repos.
func TestHFPage_LLMTabDropsOtherTasks(t *testing.T) {
raw := []Model{
{ID: "unsloth/Qwen3-Coder-30B-A3B-Instruct-GGUF", Pipeline: "text-generation"},
{ID: "unsloth/Qwen3.8-27B-GGUF", Tags: []string{"gguf", "conversational"}},
{ID: "poolside/Laguna-S-2.1-GGUF", Tags: []string{"gguf"}},
{ID: "unsloth/Qwen3.6-27B-GGUF", Pipeline: "image-text-to-text"},
{ID: "google/gemma-4-12B-it-qat-q4_0-gguf", Pipeline: "any-to-any"},
{ID: "audio-cpp/audio.cpp-gguf", Pipeline: "text-to-speech"},
{ID: "handy-computer/parakeet-tdt-0.6b-v3-gguf", Pipeline: "automatic-speech-recognition"},
{ID: "QuantStack/Wan2.2-T2V-A14B-GGUF", Pipeline: "text-to-video"},
{ID: "mixedbread-ai/mxbai-embed-large-v1", Pipeline: "feature-extraction"},
{ID: "untagged/video", Tags: []string{"gguf", "image-to-video"}},
}
want := []string{
"unsloth/Qwen3-Coder-30B-A3B-Instruct-GGUF",
"unsloth/Qwen3.8-27B-GGUF",
"poolside/Laguna-S-2.1-GGUF",
"unsloth/Qwen3.6-27B-GGUF",
"google/gemma-4-12B-it-qat-q4_0-gguf",
}
for _, kind := range []string{"llm", ""} {
p := hfPage(raw, Query{Kind: kind}, len(raw), len(raw))
var got []string
for _, m := range p.Models {
got = append(got, m.ID)
}
if strings.Join(got, ",") != strings.Join(want, ",") {
t.Errorf("kind %q: got %v, want %v", kind, got, want)
}
// The offset counts hub rows, not survivors.
if p.NextSkip != len(raw) {
t.Errorf("kind %q: NextSkip %d, want %d", kind, p.NextSkip, len(raw))
}
}
// Other tabs are the hub's own verdict and pass through untouched.
if p := hfPage(raw, Query{Kind: "any"}, len(raw), len(raw)); len(p.Models) != len(raw) {
t.Errorf("kind any: %d rows, want %d", len(p.Models), len(raw))
}
}