661 lines
25 KiB
Python
661 lines
25 KiB
Python
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()
|