use std::fs;
use std::path::Path;
use std::path::PathBuf;
use tempfile::tempdir;
use super::ID;
use super::REPO;
use super::SHA;
use super::body;
use super::config;
use super::hub;
use super::layout::blob_name;
use super::server::Reply;
use super::server::Server;
use super::server::closed_url;
use crate::error::CanddingError;
use crate::hub::cache;
use crate::hub::download::fetch;
use crate::hub::resolve_hub;
use crate::registry::WeightsFormat;
fn refusing(file: &'static str) -> impl Fn(&str, &str) -> Reply + Send + Sync + 'static {
move |method, path| match method == "HEAD" && path.ends_with(&format!("/resolve/{SHA}/{file}")) {
true => Reply::new(403, ""),
false => hub(method, path),
}
}
#[test]
fn an_interrupted_branch_download_writes_no_ref_and_resolves_again() {
let server = Server::start(refusing("config.json"));
let cache = tempdir().unwrap();
let config = config(cache.path(), &server.url);
assert!(fetch(&config, ID, "main", WeightsFormat::Safetensors).is_err());
let repo = cache.path().join(REPO);
let first = repo
.join("snapshots")
.join(SHA)
.join("1_Pooling/config.json");
assert!(first.is_file(), "the first file arrived");
assert!(!repo.join("refs/main").exists());
let before = server.requests().len();
assert!(resolve_hub(&config, ID, "main", WeightsFormat::Safetensors).is_err());
assert!(server.requests().len() > before, "taken as complete");
}
#[test]
fn an_interrupted_pinned_download_is_not_taken_as_complete() {
let server = Server::start(refusing("config.json"));
let cache = tempdir().unwrap();
let config = config(cache.path(), &server.url);
assert!(fetch(&config, ID, SHA, WeightsFormat::Safetensors).is_err());
assert!(cache::in_progress_marker(cache.path(), ID, SHA).is_file());
let before = server.requests().len();
assert!(resolve_hub(&config, ID, SHA, WeightsFormat::Safetensors).is_err());
assert!(server.requests().len() > before, "taken as complete");
}
#[test]
fn an_interrupted_repair_keeps_a_branch_from_resolving_offline() {
let cache = tempdir().unwrap();
let complete = Server::start(hub);
let mut config = config(cache.path(), &complete.url);
let snapshot = fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap();
fs::remove_file(snapshot.join("model.safetensors")).unwrap();
config.endpoint = Some(Server::start(refusing("model.safetensors")).url);
assert!(fetch(&config, ID, "main", WeightsFormat::Safetensors).is_err());
config.offline = true;
match resolve_hub(&config, ID, "main", WeightsFormat::Safetensors) {
Err(CanddingError::Offline { id }) => assert_eq!(id, ID),
other => panic!("{other:?}"),
}
}
#[test]
fn a_completed_pinned_download_resolves_without_a_request() {
let server = Server::start(hub);
let cache = tempdir().unwrap();
let config = config(cache.path(), &server.url);
let snapshot = fetch(&config, ID, SHA, WeightsFormat::Safetensors).unwrap();
assert!(!cache::in_progress_marker(cache.path(), ID, SHA).exists());
assert!(!cache.path().join(REPO).join("refs").exists());
let before = server.requests().len();
let resolved = resolve_hub(&config, ID, SHA, WeightsFormat::Safetensors).unwrap();
assert_eq!(resolved, snapshot);
assert_eq!(server.requests().len(), before);
}
#[cfg(unix)]
fn another_tool_s_pinned_snapshot(cache: &Path) -> PathBuf {
let repo = cache.join(REPO);
let snapshot = repo.join("snapshots").join(SHA);
for file in ["1_Pooling/config.json", "config.json", "model.safetensors"] {
let name = blob_name(file);
fs::create_dir_all(repo.join("blobs")).unwrap();
fs::write(repo.join("blobs").join(&name), body(file)).unwrap();
let pointer = snapshot.join(file);
fs::create_dir_all(pointer.parent().unwrap()).unwrap();
let up = "../".repeat(file.matches('/').count() + 2);
std::os::unix::fs::symlink(format!("{up}blobs/{name}"), pointer).unwrap();
}
fs::create_dir_all(cache.join(".locks").join(REPO)).unwrap();
snapshot
}
#[cfg(unix)]
#[test]
fn a_pinned_snapshot_another_tool_wrote_resolves_offline_unless_marked_in_progress() {
let cache = tempdir().unwrap();
let snapshot = another_tool_s_pinned_snapshot(cache.path());
let mut config = config(cache.path(), &closed_url());
config.offline = true;
let resolved = resolve_hub(&config, ID, SHA, WeightsFormat::Safetensors).unwrap();
assert_eq!(resolved, snapshot);
let marker = cache::in_progress_marker(cache.path(), ID, SHA);
fs::create_dir_all(marker.parent().unwrap()).unwrap();
fs::write(&marker, "").unwrap();
match resolve_hub(&config, ID, SHA, WeightsFormat::Safetensors) {
Err(CanddingError::Offline { id }) => assert_eq!(id, ID),
other => panic!("{other:?}"),
}
}