use std::collections::HashMap;
use std::path::Path;
use modelshelf::gguf::fixtures::{gguf_bytes, write_gguf};
use modelshelf::scan::fixtures;
use modelshelf::{Ecosystem, ModelQuery, OpenOptions, ScanEnv, ScanOptions, Shelf};
struct Fixture {
_tmp: tempfile::TempDir,
shelf: Shelf,
options: ScanOptions,
ollama_root: std::path::PathBuf,
lmstudio_root: std::path::PathBuf,
hf_root: std::path::PathBuf,
}
fn fixture() -> Fixture {
let tmp = tempfile::tempdir().unwrap();
let base = tmp.path();
let shelf_home = base.join("shelf");
let fake_home = base.join("home"); let ollama_root = base.join("ollama");
let lmstudio_root = base.join("lmstudio");
let hf_root = base.join("hf");
for d in [&fake_home, &ollama_root, &lmstudio_root, &hf_root] {
std::fs::create_dir_all(d).unwrap();
}
let shelf = Shelf::open(OpenOptions::at(&shelf_home).app_id("com.example.test")).unwrap();
let options = ScanOptions {
env: ScanEnv {
home_dir: Some(fake_home),
roots: HashMap::from([
(Ecosystem::Ollama, ollama_root.clone()),
(Ecosystem::Lmstudio, lmstudio_root.clone()),
(Ecosystem::HfCache, hf_root.clone()),
]),
},
..Default::default()
};
Fixture {
_tmp: tmp,
shelf,
options,
ollama_root,
lmstudio_root,
hf_root,
}
}
fn lmstudio_path(root: &Path, repo: &str, file: &str) -> std::path::PathBuf {
root.join(repo).join(file)
}
#[test]
fn scan_unifies_models_across_ecosystems_and_detects_duplicates() {
let fx = fixture();
let shared = gguf_bytes("Shared Llama", "llama", 15, 8192, b"SHARED-CONTENT");
let shared_hex =
fixtures::ollama::add_model(&fx.ollama_root, "shared-llama", "latest", &shared);
write_gguf(
&lmstudio_path(
&fx.lmstudio_root,
"bartowski/Shared-Llama-GGUF",
"Shared-Llama-Q4_K_M.gguf",
),
"Shared Llama",
"llama",
15,
8192,
b"SHARED-CONTENT",
);
fixtures::hf_cache::add_model(
&fx.hf_root,
"bartowski/Shared-Llama-GGUF",
"deadbeefrev",
"Shared-Llama-Q4_K_M.gguf",
&shared,
);
write_gguf(
&lmstudio_path(&fx.lmstudio_root, "Qwen/Solo-GGUF", "Solo-Q8_0.gguf"),
"Solo Qwen",
"qwen2",
7,
32768,
b"SOLO-CONTENT",
);
let report = fx.shelf.scan(&fx.options).unwrap();
assert_eq!(report.files_found, 4, "3 shared copies + 1 solo");
assert_eq!(report.new_models, 2, "same content collapses to one model");
assert!(report.skipped.is_empty(), "skipped: {:?}", report.skipped);
let models = fx.shelf.list().unwrap();
assert_eq!(models.len(), 2);
let shared_model = models
.iter()
.find(|m| m.id.sha256_hex() == Some(shared_hex.as_str()))
.expect("shared model registered under its content hash");
assert_eq!(shared_model.locations.len(), 3);
let ecosystems: Vec<Ecosystem> = shared_model.locations.iter().map(|l| l.ecosystem).collect();
assert!(ecosystems.contains(&Ecosystem::Ollama));
assert!(ecosystems.contains(&Ecosystem::Lmstudio));
assert!(ecosystems.contains(&Ecosystem::HfCache));
assert_eq!(
shared_model.source.as_ref().map(|s| s.repo.as_str()),
Some("bartowski/Shared-Llama-GGUF")
);
assert_eq!(report.duplicate_groups.len(), 1);
let dup = &report.duplicate_groups[0];
assert_eq!(dup.id, shared_model.id);
assert_eq!(dup.distinct_copies, 3);
assert_eq!(
dup.reclaimable_bytes,
shared_model.size_bytes * 2,
"two of three copies are reclaimable"
);
let solo = models.iter().find(|m| m.id != shared_model.id).unwrap();
let gguf = solo.gguf.as_ref().expect("gguf metadata parsed");
assert_eq!(gguf.general_name.as_deref(), Some("Solo Qwen"));
assert_eq!(gguf.quantization.as_deref(), Some("Q8_0"));
}
#[test]
fn rescan_is_idempotent_and_prunes_deleted_files() {
let fx = fixture();
let solo_path = lmstudio_path(&fx.lmstudio_root, "org/Repo-GGUF", "model-Q4_0.gguf");
write_gguf(&solo_path, "Prunable", "llama", 2, 2048, b"PRUNE-ME");
write_gguf(
&lmstudio_path(&fx.lmstudio_root, "org/Keeper-GGUF", "keeper-Q4_0.gguf"),
"Keeper",
"llama",
2,
2048,
b"KEEP-ME",
);
let first = fx.shelf.scan(&fx.options).unwrap();
assert_eq!(first.new_models, 2);
let second = fx.shelf.scan(&fx.options).unwrap();
assert_eq!(second.new_models, 0, "rescan must not re-register models");
assert_eq!(fx.shelf.list().unwrap().len(), 2);
std::fs::remove_file(&solo_path).unwrap();
let third = fx.shelf.scan(&fx.options).unwrap();
assert_eq!(third.locations_pruned, 1);
let models = fx.shelf.list().unwrap();
assert_eq!(models.len(), 1);
assert_eq!(models[0].display_name, "Keeper");
}
#[test]
fn registry_survives_reopen_and_supports_find_and_refs() {
let fx = fixture();
write_gguf(
&lmstudio_path(
&fx.lmstudio_root,
"meta/Llama-3.1-GGUF",
"Llama-3.1-Q4_K_M.gguf",
),
"Meta Llama 3.1 8B Instruct",
"llama",
15,
131072,
b"LLAMA31",
);
fx.shelf.scan(&fx.options).unwrap();
let reopened =
Shelf::open(OpenOptions::at(fx.shelf.paths().home()).app_id("com.example.other")).unwrap();
let hits = reopened
.find(&ModelQuery::name_contains("llama-3.1").quant("Q4_K_M"))
.unwrap();
assert_eq!(hits.len(), 1);
let id = hits[0].id.clone();
reopened.add_ref(&id, "summarizer").unwrap();
let m = &reopened
.find(&ModelQuery::name_contains("llama-3.1"))
.unwrap()[0];
assert_eq!(m.refs.len(), 1);
assert_eq!(m.refs[0].app_id, "com.example.other");
reopened.remove_ref(&id).unwrap();
let m = &reopened
.find(&ModelQuery::name_contains("llama-3.1"))
.unwrap()[0];
assert!(m.refs.is_empty());
}
#[test]
fn custom_roots_are_scanned() {
let fx = fixture();
let custom = fx.shelf.paths().home().parent().unwrap().join("my-models");
write_gguf(
&custom.join("nested/dir/custom-model.gguf"),
"Custom Model",
"phi3",
18,
4096,
b"CUSTOM",
);
let mut options = fx.options.clone();
options.extra_roots.push(custom);
let report = fx.shelf.scan(&options).unwrap();
assert_eq!(report.new_models, 1);
let models = fx.shelf.list().unwrap();
assert_eq!(models[0].locations[0].ecosystem, Ecosystem::Custom);
}