#[test]
fn test_extract_hf_repo_trailing_slash_after_repo() {
assert_eq!(
extract_hf_repo("hf://org/repo/"),
Some("org/repo".to_string())
);
}
#[test]
fn test_extract_hf_repo_with_multiple_trailing_slashes() {
let uri = "hf://org/repo///";
assert_eq!(extract_hf_repo(uri), Some("org/repo".to_string()));
}
#[test]
fn test_extract_hf_repo_just_hf_no_colon_slash() {
assert_eq!(extract_hf_repo("hf"), None);
}
#[test]
fn test_extract_hf_repo_hf_colon_no_slashes() {
assert_eq!(extract_hf_repo("hf:org/repo"), None);
}
#[test]
fn test_extract_hf_repo_hf_single_slash() {
assert_eq!(extract_hf_repo("hf:/org/repo"), None);
}
#[test]
fn test_extract_hf_repo_numeric_org_and_repo() {
let uri = "hf://12345/67890/model.safetensors";
assert_eq!(extract_hf_repo(uri), Some("12345/67890".to_string()));
}
#[test]
fn test_extract_hf_repo_hyphenated_names() {
let uri = "hf://my-org/my-awesome-model-v2/weights.safetensors";
assert_eq!(
extract_hf_repo(uri),
Some("my-org/my-awesome-model-v2".to_string())
);
}
#[test]
fn test_extract_hf_repo_underscored_names() {
let uri = "hf://my_org/my_model_v2";
assert_eq!(extract_hf_repo(uri), Some("my_org/my_model_v2".to_string()));
}
#[test]
fn test_extract_hf_repo_very_long_names() {
let long_org = "a".repeat(100);
let long_repo = "b".repeat(200);
let uri = format!("hf://{}/{}/model.safetensors", long_org, long_repo);
assert_eq!(
extract_hf_repo(&uri),
Some(format!("{}/{}", long_org, long_repo))
);
}
#[test]
fn test_extract_hf_repo_with_at_symbol() {
let uri = "hf://org/repo@main/model.safetensors";
assert_eq!(extract_hf_repo(uri), Some("org/repo@main".to_string()));
}
#[test]
fn test_extract_hf_repo_empty_after_prefix() {
assert_eq!(extract_hf_repo("hf://"), None);
}
#[test]
fn test_extract_hf_repo_single_char_org_and_repo() {
assert_eq!(extract_hf_repo("hf://a/b"), Some("a/b".to_string()));
}
#[test]
fn test_extract_shard_files_truncated_json() {
let json = r#"{"weight_map": {"a.weight": "model-00001.safe"#;
let shards = extract_shard_files_from_index(json);
assert!(
shards.is_empty(),
"Truncated JSON should produce no results"
);
}
#[test]
fn test_extract_shard_files_unicode_tensor_names() {
let json = r#"{
"weight_map": {
"模型.层.0.权重": "model-00001-of-00001.safetensors"
}
}"#;
let shards = extract_shard_files_from_index(json);
assert_eq!(shards.len(), 1);
assert_eq!(shards[0], "model-00001-of-00001.safetensors");
}
#[test]
fn test_extract_shard_files_colons_in_tensor_names() {
let json = r#"{
"weight_map": {
"model:layers:0:weight": "shard-001.safetensors",
"model:layers:1:weight": "shard-002.safetensors"
}
}"#;
let shards = extract_shard_files_from_index(json);
assert_eq!(shards.len(), 2);
}
#[test]
fn test_extract_shard_files_empty_json_object() {
let json = "{}";
let shards = extract_shard_files_from_index(json);
assert!(shards.is_empty());
}
#[test]
fn test_extract_shard_files_null_json() {
let json = "null";
let shards = extract_shard_files_from_index(json);
assert!(shards.is_empty());
}
#[test]
fn test_extract_shard_files_array_instead_of_object() {
let json = r#"[{"weight_map": {"a": "model.safetensors"}}]"#;
let shards = extract_shard_files_from_index(json);
assert_eq!(shards.len(), 1);
}
#[test]
fn test_extract_shard_files_weight_map_with_empty_value() {
let json = r#"{"weight_map": {"a.weight": ""}}"#;
let shards = extract_shard_files_from_index(json);
assert!(shards.is_empty(), "Empty filename should be excluded");
}
#[test]
fn test_extract_shard_files_weight_map_value_not_safetensors() {
let json = r#"{"weight_map": {"a.weight": "model.gguf"}}"#;
let shards = extract_shard_files_from_index(json);
assert!(
shards.is_empty(),
"Non-safetensors files should be excluded"
);
}
#[test]
fn test_extract_shard_files_large_model_40_shards() {
let mut entries = Vec::new();
for i in 1..=200 {
let shard_num = (i % 40) + 1;
entries.push(format!(
"\"tensor_{}\": \"model-{:05}-of-00040.safetensors\"",
i, shard_num
));
}
let json = format!("{{\"weight_map\": {{{}}}}}", entries.join(",\n"));
let shards = extract_shard_files_from_index(&json);
assert_eq!(shards.len(), 40);
assert_eq!(shards[0], "model-00001-of-00040.safetensors");
assert_eq!(shards[39], "model-00040-of-00040.safetensors");
}
#[test]
fn test_extract_shard_files_weight_map_appears_in_metadata() {
let json = r#"{
"metadata": {"description": "This model has a weight_map section"},
"weight_map": {
"a.weight": "actual-shard.safetensors"
}
}"#;
let shards = extract_shard_files_from_index(json);
assert!(
!shards.is_empty(),
"Should find shards from the actual weight_map"
);
}
#[test]
fn test_extract_shard_files_sorted_alphanumerically() {
let json = r#"{
"weight_map": {
"z": "shard-c.safetensors",
"y": "shard-a.safetensors",
"x": "shard-b.safetensors"
}
}"#;
let shards = extract_shard_files_from_index(json);
assert_eq!(shards.len(), 3);
assert_eq!(shards[0], "shard-a.safetensors");
assert_eq!(shards[1], "shard-b.safetensors");
assert_eq!(shards[2], "shard-c.safetensors");
}
#[test]
fn test_extract_shard_files_weight_map_no_opening_brace() {
let json = r#"{"weight_map": "just a string, no object here"}"#;
let shards = extract_shard_files_from_index(json);
assert!(shards.is_empty());
}
#[test]
fn test_extract_shard_files_multiple_weight_map_keys() {
let json = r#"{
"weight_map": {"a": "first.safetensors"},
"weight_map": {"b": "second.safetensors"}
}"#;
let shards = extract_shard_files_from_index(json);
assert_eq!(shards.len(), 1);
assert_eq!(shards[0], "first.safetensors");
}
#[test]
fn test_resolved_model_single_file_debug() {
let model = ResolvedModel::SingleFile("test.gguf".to_string());
let debug = format!("{:?}", model);
assert!(debug.contains("SingleFile"));
assert!(debug.contains("test.gguf"));
}
#[test]
fn test_resolved_model_sharded_debug() {
let model = ResolvedModel::Sharded {
org: "Qwen".to_string(),
repo: "Qwen2.5-Coder-3B".to_string(),
shard_files: vec!["shard-001.safetensors".to_string()],
};
let debug = format!("{:?}", model);
assert!(debug.contains("Sharded"));
assert!(debug.contains("Qwen"));
assert!(debug.contains("Qwen2.5-Coder-3B"));
assert!(debug.contains("shard-001.safetensors"));
}
#[test]
fn test_resolved_model_sharded_empty_shard_files() {
let model = ResolvedModel::Sharded {
org: "org".to_string(),
repo: "repo".to_string(),
shard_files: vec![],
};
match model {
ResolvedModel::Sharded { shard_files, .. } => {
assert!(shard_files.is_empty());
}
_ => panic!("Expected Sharded"),
}
}
#[test]
fn test_shard_manifest_deserialize_unknown_fields() {
let json = r#"{"version": 2, "repo": "org/repo", "files": {}, "extra_field": "ignored"}"#;
let manifest: ShardManifest =
serde_json::from_str(json).expect("unknown fields should be ignored by default");
assert_eq!(manifest.version, 2);
}
#[test]
fn test_shard_manifest_missing_required_field() {
let json = r#"{"version": 1, "files": {}}"#;
assert!(serde_json::from_str::<ShardManifest>(json).is_err());
}
#[test]
fn test_file_checksum_missing_blake3_field() {
let json = r#"{"size": 100}"#;
assert!(serde_json::from_str::<FileChecksum>(json).is_err());
}
#[test]
fn test_file_checksum_missing_size_field() {
let json = r#"{"blake3": "abc123"}"#;
assert!(serde_json::from_str::<FileChecksum>(json).is_err());
}
#[test]
fn test_shard_manifest_special_chars_in_repo() {
let manifest = ShardManifest {
version: 1,
repo: "org/repo-with.dots_and-dashes".to_string(),
files: HashMap::new(),
};
let json = serde_json::to_string(&manifest).expect("serialize");
let deserialized: ShardManifest = serde_json::from_str(&json).expect("deserialize");
assert_eq!(deserialized.repo, "org/repo-with.dots_and-dashes");
}
#[test]
fn test_shard_manifest_many_files() {
let mut files = HashMap::new();
for i in 0..100 {
files.insert(
format!("model-{:05}-of-00100.safetensors", i + 1),
FileChecksum {
size: 5_000_000_000 + i as u64,
blake3: format!("hash_{:05}", i),
},
);
}
let manifest = ShardManifest {
version: 1,
repo: "big/model".to_string(),
files,
};
let json = serde_json::to_string(&manifest).expect("serialize");
let deserialized: ShardManifest = serde_json::from_str(&json).expect("deserialize");
assert_eq!(deserialized.files.len(), 100);
}
#[test]
fn test_file_checksum_unicode_blake3() {
let checksum = FileChecksum {
size: 42,
blake3: "hash_with_\u{00e9}moji".to_string(),
};
let json = serde_json::to_string(&checksum).expect("serialize");
let deserialized: FileChecksum = serde_json::from_str(&json).expect("deserialize");
assert_eq!(deserialized.blake3, "hash_with_\u{00e9}moji");
}
#[test]
fn test_resolve_hf_model_bare_org_repo_with_pt_extension() {
let result = resolve_hf_model("org/repo/model.pt").expect("pt'");
match result {
ResolvedModel::SingleFile(s) => assert_eq!(s, "hf://org/repo/model.pt"),
ResolvedModel::Sharded { .. } => panic!("Expected SingleFile"),
}
}
#[test]
fn test_resolve_hf_model_file_url_scheme() {
let result = resolve_hf_model("file:///path/to/model.gguf").expect("gguf'");
match result {
ResolvedModel::SingleFile(s) => assert_eq!(s, "file:///path/to/model.gguf"),
ResolvedModel::Sharded { .. } => panic!("Expected SingleFile"),
}
}
#[test]
fn test_resolve_hf_model_s3_url_scheme() {
let result = resolve_hf_model("s3://bucket/model.gguf").expect("gguf'");
match result {
ResolvedModel::SingleFile(s) => assert_eq!(s, "s3://bucket/model.gguf"),
ResolvedModel::Sharded { .. } => panic!("Expected SingleFile"),
}
}
#[test]
fn test_resolve_hf_model_just_dot() {
let result = resolve_hf_model(".").expect("'");
match result {
ResolvedModel::SingleFile(s) => assert_eq!(s, "."),
ResolvedModel::Sharded { .. } => panic!("Expected SingleFile"),
}
}
#[test]
fn test_resolve_hf_model_empty_string() {
let result = resolve_hf_model("").expect("resolve_hf_model('')");
match result {
ResolvedModel::SingleFile(s) => assert_eq!(s, ""),
ResolvedModel::Sharded { .. } => panic!("Expected SingleFile"),
}
}
#[test]
fn test_resolve_hf_model_gguf_mixed_case_extension() {
let result = resolve_hf_model("hf://org/repo/model.GGuF").expect("GGuF'");
match result {
ResolvedModel::SingleFile(s) => assert_eq!(s, "hf://org/repo/model.GGuF"),
ResolvedModel::Sharded { .. } => panic!("Expected SingleFile for .GGuF"),
}
}
#[test]
fn test_resolve_hf_model_pt_mixed_case() {
let result = resolve_hf_model("hf://org/repo/model.PT").expect("PT'");
match result {
ResolvedModel::SingleFile(s) => assert_eq!(s, "hf://org/repo/model.PT"),
ResolvedModel::Sharded { .. } => panic!("Expected SingleFile"),
}
}
#[test]
fn test_resolve_hf_uri_single_file_with_safetensors() {
let uri = "hf://org/repo/model.safetensors";
let resolved = resolve_hf_uri(uri).expect("safetensors'; let reso");
assert_eq!(resolved, uri);
}