import hashlib import json import os import tempfile import unittest from unittest.mock import patch from core import worker_utils from core.metadata_builder import ( METADATA_BUILD_MODE_CALCULATE_FRESH, METADATA_BUILD_MODE_IMPORT_EXISTING, _empty_build_counts, build_missing_local_metadata, normalize_metadata_build_mode, ) from core.metadata_utils import ( build_metadata_source_payload, build_sidecar_file_identity, merge_counted_payload, ) from core.path_utils import get_model_resolver_sidecar_path class MetadataBuilderTests(unittest.TestCase): def test_build_metadata_source_payload_preserves_common_source_fields(self): payload = build_metadata_source_payload( source="civitai", details_source="civitai", filename="model.safetensors", category="checkpoints", model_name="Example Model", model_type="Checkpoint", sha256="a" * 64, size=1234, result={ "base_model": "SDXL", "base_model_source": "civitai", "base_model_inferred": 1, }, description="Model description", version_description="Version description", extra_fields={"author": "Author", "size": 5678}, ) self.assertEqual( { "source": "civitai", "details_source": "civitai", "filename": "model.safetensors", "category": "checkpoints", "model_name": "Example Model", "name": "Example Model", "model_type": "Checkpoint", "sha256": "a" * 64, "size": 5678, "base_model": "SDXL", "base_model_source": "civitai", "base_model_inferred": True, "description": "Model description", "model_description": "Model description", "version_description": "Version description", "author": "Author", }, payload, ) def test_build_sidecar_file_identity_returns_common_schema_fields(self): identity = build_sidecar_file_identity( r"models\example.safetensors", model_name="Example Model", size=1234, modified_at=10.0, last_checked_at=11.0, ) self.assertEqual( { "schema": "comfyui-model-resolver", "schema_version": 1, "managed_by": "comfyui-model-resolver", "file_name": "example", "filename": "example.safetensors", "model_name": "Example Model", "file_path": "models/example.safetensors", "size": 1234, "modified": 10.0, "last_checked_at": 11.0, }, identity, ) def test_merge_counted_payload_sums_build_numbers_and_preserves_list_order(self): target = _empty_build_counts() merge_counted_payload( target, { "scanned_models": 2, "created_metadata": 1, "errors": [{"message": "first"}], "updated": [{"filename": "first.safetensors"}], "history": [{"action": "created"}], }, numeric_keys=( "scanned_models", "created_metadata", "updated_metadata", "skipped_complete", "header_hashes", "calculated_hashes", "header_metadata_count", "invalid_metadata", ), list_keys=("errors", "updated", "history"), ) merge_counted_payload( target, { "scanned_models": 3, "created_metadata": 2, "errors": [{"message": "second"}], "updated": [{"filename": "second.safetensors"}], "history": [{"action": "updated"}], }, numeric_keys=( "scanned_models", "created_metadata", "updated_metadata", "skipped_complete", "header_hashes", "calculated_hashes", "header_metadata_count", "invalid_metadata", ), list_keys=("errors", "updated", "history"), ) self.assertEqual(5, target["scanned_models"]) self.assertEqual(3, target["created_metadata"]) self.assertEqual( [{"message": "first"}, {"message": "second"}], target["errors"], ) self.assertEqual( [ {"filename": "first.safetensors"}, {"filename": "second.safetensors"}, ], target["updated"], ) self.assertEqual( [{"action": "created"}, {"action": "updated"}], target["history"], ) def _write_model(self, directory, filename, content): path = os.path.join(directory, filename) with open(path, "wb") as handle: handle.write(content) return path def _write_fake_safetensors(self, directory, filename, metadata): header = json.dumps({"__metadata__": metadata}).encode("utf-8") content = len(header).to_bytes(8, "little") + header return self._write_model(directory, filename, content) def _write_metadata(self, model_path, payload): metadata_path = get_model_resolver_sidecar_path(model_path) with open(metadata_path, "w", encoding="utf-8") as handle: json.dump(payload, handle) return metadata_path def _read_metadata(self, model_path): metadata_path = get_model_resolver_sidecar_path(model_path) with open(metadata_path, "r", encoding="utf-8") as handle: return metadata_path, json.load(handle) def _model_info(self, model_path, category="loras"): return { "filename": os.path.basename(model_path), "relative_path": os.path.basename(model_path), "path": model_path, "category": category, "base_directory": os.path.dirname(model_path), } def test_worker_count_preserves_builder_defaults_and_requested_limits(self): with patch.object(worker_utils.os, "cpu_count", return_value=8): self.assertEqual( worker_utils.resolve_worker_count( 20, default_worker_limit=4, ), (4, 8), ) self.assertEqual( worker_utils.resolve_worker_count( 20, 6, default_worker_limit=4, ), (6, 8), ) self.assertEqual( worker_utils.resolve_worker_count( 0, default_worker_limit=4, ), (4, 8), ) def test_creates_missing_metadata_from_local_safetensors_header(self): with tempfile.TemporaryDirectory() as tmpdir: header_sha = "a" * 64 model_path = self._write_fake_safetensors( tmpdir, "header_model.safetensors", { "modelspec.title": "Header Model", "modelspec.description": "Local description", "modelspec.tags": "tag-a, tag-b", "modelspec.trigger_phrase": "trigger-a", "sha256": header_sha, }, ) result = build_missing_local_metadata( models=[self._model_info(model_path)], metadata_mode=METADATA_BUILD_MODE_IMPORT_EXISTING, ) metadata_path, metadata = self._read_metadata(model_path) self.assertTrue(result["success"]) self.assertEqual(1, result["created_metadata"]) self.assertEqual(0, result["calculated_hashes"]) self.assertEqual(1, result["header_hashes"]) self.assertEqual(metadata_path, result["updated"][0]["metadata_path"]) self.assertEqual("Header Model", metadata["model_name"]) self.assertEqual(header_sha, metadata["sha256"]) self.assertEqual("safetensors_header", metadata["sha256_source"]) self.assertIn("tag-a", metadata["tags"]) self.assertIn("trigger-a", metadata["trained_words"]) self.assertIn("safetensors_header_metadata", metadata) self.assertIn("modelspec.title", metadata["safetensors_header_metadata"]["metadata"]) def test_updates_existing_metadata_without_sha256_and_preserves_user_fields(self): with tempfile.TemporaryDirectory() as tmpdir: content = b"local model bytes" expected_sha = hashlib.sha256(content).hexdigest() model_path = self._write_model(tmpdir, "needs_hash.ckpt", content) self._write_metadata( model_path, { "model_name": "Existing Name", "notes": "keep me", "favorite": True, }, ) events = [] result = build_missing_local_metadata( models=[self._model_info(model_path, "checkpoints")], progress_callback=events.append, ) _metadata_path, metadata = self._read_metadata(model_path) self.assertTrue(result["success"]) self.assertEqual(1, result["updated_metadata"]) self.assertEqual(1, result["calculated_hashes"]) self.assertEqual(expected_sha, metadata["sha256"]) self.assertEqual("file", metadata["sha256_source"]) self.assertEqual("Existing Name", metadata["model_name"]) self.assertEqual("keep me", metadata["notes"]) self.assertTrue(metadata["favorite"]) self.assertTrue(any(event.get("stage") == "hashing" for event in events)) self.assertTrue(any(event.get("current_model") == "needs_hash.ckpt" for event in events)) def test_skips_existing_metadata_that_already_has_sha256(self): with tempfile.TemporaryDirectory() as tmpdir: model_path = self._write_model(tmpdir, "complete.safetensors", b"abc") existing = { "model_name": "Complete", "sha256": "b" * 64, "notes": "do not rewrite", } self._write_metadata(model_path, existing) result = build_missing_local_metadata(models=[self._model_info(model_path)]) _metadata_path, metadata = self._read_metadata(model_path) self.assertTrue(result["success"]) self.assertEqual(0, result["created_metadata"]) self.assertEqual(0, result["updated_metadata"]) self.assertEqual(1, result["skipped_complete"]) self.assertEqual(existing, metadata) def test_same_stem_models_receive_extension_qualified_sidecars(self): with tempfile.TemporaryDirectory() as tmpdir: bin_content = b"bin adapter weights" safetensors_content = b"safetensors adapter weights" bin_path = self._write_model(tmpdir, "ip-adapter_sd15.bin", bin_content) safetensors_path = self._write_model( tmpdir, "ip-adapter_sd15.safetensors", safetensors_content, ) result = build_missing_local_metadata( models=[ self._model_info(bin_path), self._model_info(safetensors_path), ], worker_count=1, metadata_mode=METADATA_BUILD_MODE_CALCULATE_FRESH, ) bin_metadata_path = get_model_resolver_sidecar_path(bin_path) safetensors_metadata_path = get_model_resolver_sidecar_path( safetensors_path ) legacy_metadata_path = os.path.splitext(bin_path)[0] + ".metadata.json" with open(bin_metadata_path, "r", encoding="utf-8") as handle: bin_metadata = json.load(handle) with open(safetensors_metadata_path, "r", encoding="utf-8") as handle: safetensors_metadata = json.load(handle) legacy_metadata_exists = os.path.exists(legacy_metadata_path) self.assertTrue(result["success"]) self.assertEqual(2, result["created_metadata"]) self.assertFalse(legacy_metadata_exists) self.assertEqual("ip-adapter_sd15.bin", bin_metadata["filename"]) self.assertEqual( hashlib.sha256(bin_content).hexdigest(), bin_metadata["sha256"], ) self.assertEqual( "ip-adapter_sd15.safetensors", safetensors_metadata["filename"], ) self.assertEqual( hashlib.sha256(safetensors_content).hexdigest(), safetensors_metadata["sha256"], ) self.assertEqual("comfyui-model-resolver", bin_metadata["managed_by"]) self.assertNotIn("favorite", bin_metadata) self.assertNotIn("notes", bin_metadata) self.assertEqual( METADATA_BUILD_MODE_CALCULATE_FRESH, bin_metadata["metadata_build_mode"], ) def test_external_metadata_remains_untouched_and_is_not_copied(self): with tempfile.TemporaryDirectory() as tmpdir: bin_content = b"bin" bin_path = self._write_model(tmpdir, "adapter.bin", bin_content) safetensors_path = self._write_model( tmpdir, "adapter.safetensors", b"safe", ) legacy_sha = "c" * 64 legacy_path = os.path.splitext(bin_path)[0] + ".metadata.json" legacy_payload = { "filename": "adapter.bin", "file_path": bin_path, "sha256": legacy_sha, "notes": "preserve external user details", "favorite": True, } with open(legacy_path, "w", encoding="utf-8") as handle: json.dump(legacy_payload, handle) result = build_missing_local_metadata( models=[ self._model_info(bin_path), self._model_info(safetensors_path), ], worker_count=1, metadata_mode=METADATA_BUILD_MODE_CALCULATE_FRESH, ) with open( get_model_resolver_sidecar_path(bin_path), "r", encoding="utf-8", ) as handle: bin_metadata = json.load(handle) with open( get_model_resolver_sidecar_path(safetensors_path), "r", encoding="utf-8", ) as handle: safetensors_metadata = json.load(handle) with open(legacy_path, "r", encoding="utf-8") as handle: legacy_after = json.load(handle) self.assertTrue(result["success"]) self.assertEqual(legacy_payload, legacy_after) self.assertEqual( hashlib.sha256(bin_content).hexdigest(), bin_metadata["sha256"], ) self.assertNotEqual(legacy_sha, bin_metadata["sha256"]) self.assertNotIn("notes", bin_metadata) self.assertNotIn("favorite", bin_metadata) self.assertEqual("adapter.bin", bin_metadata["filename"]) self.assertEqual("adapter.safetensors", safetensors_metadata["filename"]) self.assertEqual( hashlib.sha256(b"safe").hexdigest(), safetensors_metadata["sha256"], ) def test_import_mode_uses_external_metadata_without_modifying_it(self): with tempfile.TemporaryDirectory() as tmpdir: model_path = self._write_model( tmpdir, "imported.safetensors", b"local model bytes", ) external_path = os.path.join(tmpdir, "imported.metadata.json") external_sha = "d" * 64 external_payload = { "filename": "imported.safetensors", "file_path": model_path, "model_name": "Imported model name", "base_model": "SDXL 1.0", "tags": ["external-tag"], "sha256": external_sha, "hash_status": "completed", "notes": "user-owned note", "favorite": True, } with open(external_path, "w", encoding="utf-8") as handle: json.dump(external_payload, handle) result = build_missing_local_metadata( models=[self._model_info(model_path)], worker_count=1, metadata_mode=METADATA_BUILD_MODE_IMPORT_EXISTING, ) _metadata_path, metadata = self._read_metadata(model_path) with open(external_path, "r", encoding="utf-8") as handle: external_after = json.load(handle) self.assertTrue(result["success"]) self.assertEqual(1, result["created_metadata"]) self.assertEqual(0, result["calculated_hashes"]) self.assertEqual(external_payload, external_after) self.assertEqual(external_sha, metadata["sha256"]) self.assertEqual("external_metadata", metadata["sha256_source"]) self.assertEqual("Imported model name", metadata["model_name"]) self.assertEqual("SDXL 1.0", metadata["base_model"]) self.assertEqual(["external-tag"], metadata["tags"]) self.assertNotIn("notes", metadata) self.assertNotIn("favorite", metadata) self.assertEqual( external_path.replace(os.sep, "/"), metadata["imported_from"], ) self.assertEqual( METADATA_BUILD_MODE_IMPORT_EXISTING, metadata["metadata_build_mode"], ) def test_import_mode_calculates_hash_when_external_hash_is_pending(self): with tempfile.TemporaryDirectory() as tmpdir: content = b"calculate missing external hash" model_path = self._write_model( tmpdir, "pending.safetensors", content, ) external_path = os.path.join(tmpdir, "pending.metadata.json") with open(external_path, "w", encoding="utf-8") as handle: json.dump( { "filename": "pending.safetensors", "sha256": "e" * 64, "hash_status": "pending", }, handle, ) result = build_missing_local_metadata( models=[self._model_info(model_path)], worker_count=1, metadata_mode=METADATA_BUILD_MODE_IMPORT_EXISTING, ) _metadata_path, metadata = self._read_metadata(model_path) self.assertEqual(1, result["calculated_hashes"]) self.assertEqual(hashlib.sha256(content).hexdigest(), metadata["sha256"]) self.assertEqual("file", metadata["sha256_source"]) def test_invalid_metadata_mode_falls_back_to_fresh_calculation(self): self.assertEqual( METADATA_BUILD_MODE_CALCULATE_FRESH, normalize_metadata_build_mode("unsupported"), ) def test_invalid_metadata_mode_type_is_rejected(self): with self.assertRaises(TypeError): normalize_metadata_build_mode(123) def test_exact_sidecar_remains_canonical_after_sibling_is_removed(self): with tempfile.TemporaryDirectory() as tmpdir: bin_path = self._write_model(tmpdir, "adapter.bin", b"bin") safetensors_path = self._write_model( tmpdir, "adapter.safetensors", b"safe", ) exact_metadata_path = get_model_resolver_sidecar_path(bin_path) with open(exact_metadata_path, "w", encoding="utf-8") as handle: json.dump({"filename": "adapter.bin", "size": 3}, handle) os.remove(safetensors_path) selected_path = get_model_resolver_sidecar_path(bin_path) self.assertEqual(exact_metadata_path, selected_path) def test_builds_missing_metadata_with_multiple_workers(self): with tempfile.TemporaryDirectory() as tmpdir: model_paths = [ self._write_model(tmpdir, f"parallel_{index}.ckpt", f"model-{index}".encode("utf-8")) for index in range(4) ] events = [] result = build_missing_local_metadata( models=[self._model_info(path, "checkpoints") for path in model_paths], worker_count=2, progress_callback=events.append, ) metadata_payloads = [self._read_metadata(path)[1] for path in model_paths] self.assertTrue(result["success"]) self.assertEqual(2, result["worker_count"]) self.assertGreaterEqual(result["cpu_count"], 1) self.assertEqual(4, result["created_metadata"]) self.assertEqual(4, result["calculated_hashes"]) self.assertEqual(4, len(metadata_payloads)) self.assertTrue(all(payload.get("sha256") for payload in metadata_payloads)) self.assertTrue(any(event.get("worker_count") == 2 for event in events)) def test_final_result_contract_matches_between_sequential_and_parallel_builds(self): with tempfile.TemporaryDirectory() as tmpdir: sequential_path = self._write_model( tmpdir, "sequential.ckpt", b"sequential model", ) parallel_paths = [ self._write_model( tmpdir, f"parallel_contract_{index}.ckpt", f"parallel model {index}".encode(), ) for index in range(2) ] sequential_result = build_missing_local_metadata( models=[self._model_info(sequential_path, "checkpoints")], worker_count=1, metadata_mode=METADATA_BUILD_MODE_IMPORT_EXISTING, ) parallel_result = build_missing_local_metadata( models=[ self._model_info(path, "checkpoints") for path in parallel_paths ], worker_count=2, metadata_mode=METADATA_BUILD_MODE_IMPORT_EXISTING, ) expected_keys = { "success", "cancelled", "scanned_models", "total_models", "created_metadata", "updated_metadata", "skipped_complete", "header_hashes", "calculated_hashes", "header_metadata_count", "invalid_metadata", "worker_count", "cpu_count", "error_count", "errors", "updated", "updated_count", "history", "history_count", "metadata_mode", } self.assertEqual(expected_keys, set(sequential_result)) self.assertEqual(expected_keys, set(parallel_result)) self.assertEqual( METADATA_BUILD_MODE_IMPORT_EXISTING, sequential_result["metadata_mode"], ) self.assertEqual( METADATA_BUILD_MODE_IMPORT_EXISTING, parallel_result["metadata_mode"], ) self.assertEqual( sequential_result["history_count"], len(sequential_result["history"]), ) self.assertEqual( parallel_result["history_count"], len(parallel_result["history"]), ) def test_history_items_preserve_model_context_for_created_and_skipped_files(self): with tempfile.TemporaryDirectory() as tmpdir: created_path = self._write_model( tmpdir, "history_created.ckpt", b"created model", ) skipped_path = self._write_model( tmpdir, "history_skipped.ckpt", b"skipped model", ) self._write_metadata( skipped_path, {"filename": "history_skipped.ckpt", "sha256": "f" * 64}, ) result = build_missing_local_metadata( models=[ self._model_info(created_path, "checkpoints"), self._model_info(skipped_path, "loras"), ], worker_count=1, ) history_by_filename = { item["filename"]: item for item in result["history"] } self.assertEqual("created", history_by_filename["history_created.ckpt"]["action"]) self.assertEqual("skipped", history_by_filename["history_skipped.ckpt"]["action"]) for filename, category in ( ("history_created.ckpt", "checkpoints"), ("history_skipped.ckpt", "loras"), ): with self.subTest(filename=filename): item = history_by_filename[filename] self.assertEqual(filename, item["relative_path"]) self.assertEqual(category, item["category"]) self.assertTrue(item["model_path"]) self.assertTrue(item["metadata_path"]) if __name__ == "__main__": unittest.main()