mod client;
mod complete;
mod guards;
mod headers;
mod layout;
mod listing;
mod lock;
mod python;
mod server;
mod verify;
use std::fs;
use std::path::Path;
use serde_json::json;
use server::Reply;
use server::Server;
use server::closed_url;
use tempfile::tempdir;
use super::failure::Kind;
use super::failure::Step;
use super::failure::kind;
use super::fetch;
use super::valid_repo_id;
use crate::error::CanddingError;
use crate::hub::HubConfig;
use crate::hub::resolve_hub;
use crate::registry;
use crate::registry::WeightsFormat;
const ID: &str = "org/model";
const REPO: &str = "models--org--model";
const SHA: &str = "0123456789abcdef0123456789abcdef01234567";
const FILES: [(&str, &str); 5] = [
(
"1_Pooling/config.json",
"398ac2287d9d5320d12669e75a051eda64f8e0ed",
),
("README.md", "371a0e2b64a51e5d651ca3b35a7e53a82e6e281e"),
("config.json", "9238dafaefa3fd8678e426a879234453e4d404c3"),
("model.safetensors", "an LFS pointer"),
(
"onnx/model.onnx",
"3f83b4e76ca1a9f252fc76a7102445c7e77330b5",
),
];
fn body(file: &str) -> String {
format!("bytes of {file}")
}
fn oid(file: &str) -> &'static str {
FILES
.iter()
.find(|(listed, _)| *listed == file)
.map_or("f", |(_, oid)| oid)
}
fn hub(method: &str, path: &str) -> Reply {
let path = path.split('?').next().unwrap_or(path);
let revision = |name: &str| path == format!("/api/models/{ID}/revision/{name}");
if method == "GET" && (revision("main") || revision(SHA)) {
return Reply::new(200, json!({ "id": ID, "sha": SHA }).to_string());
}
if method == "GET" && path == format!("/api/models/{ID}/tree/{SHA}") {
let mut entries = vec![json!({ "type": "directory", "oid": "d", "size": 0, "path": "onnx" })];
for (file, oid) in FILES {
let mut entry = json!({ "type": "file", "oid": oid, "size": body(file).len(), "path": file });
if file.ends_with(".safetensors") {
entry["lfs"] = json!({ "oid": "a sha256", "size": body(file).len(), "pointerSize": 134 });
}
entries.push(entry);
}
return Reply::new(200, json!(entries).to_string());
}
match path.strip_prefix(&format!("/{ID}/resolve/{SHA}/")) {
Some(file) => Reply::new(200, body(file))
.header("ETag", format!("\"blob-{}\"", file.replace('/', "-")))
.header("X-Repo-Commit", SHA),
None => Reply::new(404, ""),
}
}
fn config(cache: &Path, endpoint: &str) -> HubConfig {
HubConfig {
cache: cache.to_path_buf(),
offline: false,
endpoint: Some(endpoint.to_string()),
token: Some("hf_test".to_string()),
}
}
fn failure(route: impl Fn(&str, &str) -> Reply + Send + Sync + 'static) -> CanddingError {
let server = Server::start(route);
let cache = tempdir().unwrap();
let config = config(cache.path(), &server.url);
fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap_err()
}
fn files_under(dir: &Path, prefix: &str) -> Vec<String> {
let mut files = Vec::new();
for entry in fs::read_dir(dir).unwrap() {
let path = entry.unwrap().path();
let name = format!("{prefix}{}", path.file_name().unwrap().to_str().unwrap());
if path.is_dir() {
files.extend(files_under(&path, &format!("{name}/")));
} else {
files.push(name);
}
}
files.sort();
files
}
#[test]
fn a_fetch_links_only_the_selected_files_and_writes_the_ref() {
let server = Server::start(hub);
let cache = tempdir().unwrap();
let config = config(cache.path(), &server.url);
let snapshot = fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap();
let repo = cache.path().join("models--org--model");
assert_eq!(snapshot, repo.join("snapshots").join(SHA));
assert_eq!(fs::read_to_string(repo.join("refs/main")).unwrap(), SHA);
assert_eq!(
files_under(&snapshot, ""),
["1_Pooling/config.json", "config.json", "model.safetensors"]
);
let pooling = fs::read_to_string(snapshot.join("1_Pooling/config.json")).unwrap();
assert_eq!(pooling, "bytes of 1_Pooling/config.json");
let requests = server.requests();
assert!(
!requests
.iter()
.any(|r| r.contains("onnx/") || r.contains("README")),
"{requests:?}"
);
}
#[test]
fn a_cached_snapshot_resolves_without_a_request() {
let server = Server::start(hub);
let cache = tempdir().unwrap();
let repo = cache.path().join("models--org--model");
fs::create_dir_all(repo.join("refs")).unwrap();
fs::create_dir_all(repo.join("snapshots").join(SHA)).unwrap();
fs::write(repo.join("refs/main"), SHA).unwrap();
let config = config(cache.path(), &server.url);
let dir = resolve_hub(&config, ID, "main", WeightsFormat::Safetensors).unwrap();
assert_eq!(dir, repo.join("snapshots").join(SHA));
assert_eq!(server.requests(), Vec::<String>::new());
}
#[test]
fn offline_names_the_model_and_sends_nothing() {
let server = Server::start(hub);
let cache = tempdir().unwrap();
let mut config = config(cache.path(), &server.url);
config.offline = true;
match resolve_hub(&config, ID, "main", WeightsFormat::Safetensors) {
Err(error @ CanddingError::Offline { .. }) => assert!(error.to_string().contains(ID)),
other => panic!("expected Offline, got {other:?}"),
}
assert_eq!(server.requests(), Vec::<String>::new());
}
#[test]
fn a_401_or_403_on_a_file_is_the_gated_error() {
for status in [401, 403] {
for code in [Some("GatedRepo"), None] {
let error = failure(move |method, path| match path.contains("/resolve/") {
true if code.is_some() => Reply::new(status, "").header("X-Error-Code", "GatedRepo"),
true => Reply::new(status, ""),
false => hub(method, path),
});
assert!(
matches!(&error, CanddingError::GatedModel { id } if id == ID),
"{status} {code:?}: {error:?}"
);
}
}
}
#[test]
fn an_anonymous_401_on_the_listing_is_a_missing_repository() {
let error = failure(|_, _| {
Reply::new(401, r#"{"error":"Invalid username or password."}"#)
.header("X-Error-Message", "Invalid username or password.")
});
assert!(
matches!(&error, CanddingError::RepositoryNotFound { id } if id == ID),
"{error:?}"
);
}
#[test]
fn the_hub_s_error_code_names_a_missing_repository_or_revision() {
let error = failure(|_, _| Reply::new(404, "").header("X-Error-Code", "RepoNotFound"));
assert!(
matches!(&error, CanddingError::RepositoryNotFound { id } if id == ID),
"{error:?}"
);
let error = failure(|_, _| Reply::new(404, "").header("X-Error-Code", "RevisionNotFound"));
assert!(
matches!(&error, CanddingError::RevisionNotFound { id, revision } if id == ID && revision == "main"),
"{error:?}"
);
}
#[test]
fn an_unreachable_hub_is_a_network_error() {
let cache = tempdir().unwrap();
let config = config(cache.path(), &closed_url());
let error = fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap_err();
assert!(
matches!(&error, CanddingError::Network { id, .. } if id == ID),
"{error:?}"
);
}
#[test]
fn the_error_code_decides_before_the_status_and_the_step() {
assert_eq!(kind(Step::Listing, Some(401), None), Kind::Repository);
assert_eq!(kind(Step::Listing, Some(404), None), Kind::Repository);
assert_eq!(kind(Step::Listing, Some(403), None), Kind::Gated);
assert_eq!(kind(Step::File, Some(401), None), Kind::Gated);
assert_eq!(kind(Step::File, Some(403), None), Kind::Gated);
let revision = Some("RevisionNotFound");
assert_eq!(kind(Step::Listing, Some(404), revision), Kind::Revision);
assert_eq!(
kind(Step::File, Some(401), Some("RepoNotFound")),
Kind::Repository
);
assert_eq!(
kind(Step::Listing, Some(401), Some("GatedRepo")),
Kind::Gated
);
assert_eq!(
kind(Step::File, Some(404), Some("EntryNotFound")),
Kind::Other
);
assert_eq!(kind(Step::File, Some(503), None), Kind::Other);
assert_eq!(kind(Step::Listing, None, None), Kind::Other);
}
#[test]
fn an_id_huggingface_hub_would_refuse_never_reaches_the_hub() {
let server = Server::start(hub);
let cache = tempdir().unwrap();
let config = config(cache.path(), &server.url);
for id in [
"./bert-base-uncased",
"/abs/dir",
"../m",
"a b",
"o/n/x",
"o--x/n",
"o/n.git",
"o/-n",
"org/a..b",
] {
match fetch(&config, id, "main", WeightsFormat::Safetensors) {
Err(CanddingError::RepositoryNotFound { id: named }) => assert_eq!(named, id),
other => panic!("{id}: {other:?}"),
}
}
assert_eq!(server.requests(), Vec::<String>::new());
}
#[test]
fn every_catalog_id_is_a_valid_repository_id() {
let sparse = registry::sparse::all();
let ids = registry::dense::all()
.into_iter()
.map(|d| d.id)
.chain(sparse.iter().map(|d| d.id.clone()))
.chain(sparse.iter().filter_map(|d| d.head_source.clone()))
.chain(registry::multi_vector::all().into_iter().map(|d| d.id))
.chain(registry::rerank::all().into_iter().map(|d| d.id));
for id in ids {
assert!(valid_repo_id(&id), "{id}");
}
assert!(valid_repo_id("gpt2") && valid_repo_id("Foo-BAR_foo.bar123"));
}