mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-09 07:20:15 -03:00
feat(metadata): add CivitAI AutoV3 hash support across all storage layers
- Three-state autov3 field (not-checked / checked-unavailable / 12-hex value) in .metadata.json sidecars, in-memory ModelHashIndex, and SQLite (models.autov3 column + autov3_index table) with column-presence migration - Background self-terminating backfill for legacy rows: per-model-type concurrency guard, executor-offloaded I/O, Civitai-first resolution (SHA256-matched version file) falling back to the embedded safetensors header hash - Civitai-first propagation on metadata refresh, scan, and download paths; reject the empty-string SHA256 placeholder and strip OneTrainer 0x prefix - List API hash filters and hash index lookups accept 12-char AutoV3 - Cap safetensors header reads at 64 MiB to prevent crafted-file allocation - Prevent stale AutoV3 mappings on file replacement while preserving them on same-file re-registration (lazy-hash completion)
This commit is contained in:
Will Miao
350
tests/services/test_autov3_backfill_service.py
Normal file
350
tests/services/test_autov3_backfill_service.py
Normal file
@@ -0,0 +1,350 @@
|
||||
"""Tests for Autov3BackfillService."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from py.services.autov3_backfill_service import Autov3BackfillService
|
||||
from py.services.model_cache import ModelCache
|
||||
from py.services.model_hash_index import ModelHashIndex
|
||||
from py.services.model_scanner import ModelScanner
|
||||
from py.services.persistent_model_cache import DEFAULT_LICENSE_FLAGS, PersistentModelCache
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_backfill_singleton() -> None:
|
||||
"""Reset the service singleton so every test starts from a fresh instance."""
|
||||
Autov3BackfillService._instance = None
|
||||
yield
|
||||
Autov3BackfillService._instance = None
|
||||
|
||||
|
||||
def _entry(file_path: str, sha256: str, autov3: Optional[str] = None) -> Dict[str, Any]:
|
||||
return {
|
||||
'file_path': file_path,
|
||||
'file_name': Path(file_path).stem,
|
||||
'model_name': Path(file_path).stem,
|
||||
'folder': '',
|
||||
'size': 1,
|
||||
'modified': 1.0,
|
||||
'sha256': sha256,
|
||||
'autov3': autov3,
|
||||
'base_model': '',
|
||||
'preview_url': '',
|
||||
'preview_nsfw_level': 0,
|
||||
'from_civitai': True,
|
||||
'favorite': False,
|
||||
'notes': '',
|
||||
'usage_tips': '',
|
||||
'metadata_source': None,
|
||||
'exclude': False,
|
||||
'db_checked': False,
|
||||
'last_checked_at': 0.0,
|
||||
'tags': [],
|
||||
'civitai': None,
|
||||
'civitai_deleted': False,
|
||||
'skip_metadata_refresh': False,
|
||||
'license_flags': DEFAULT_LICENSE_FLAGS,
|
||||
'hash_status': 'completed',
|
||||
'hf_url': '',
|
||||
}
|
||||
|
||||
|
||||
class RecordingScanner:
|
||||
"""Duck-typed scanner double persisting updates to a real cache."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_type: str,
|
||||
persistent_cache: PersistentModelCache,
|
||||
entries: List[Dict[str, Any]],
|
||||
) -> None:
|
||||
self.model_type = model_type
|
||||
self._persistent_cache = persistent_cache
|
||||
self.entries: Dict[str, Dict[str, Any]] = {entry['file_path']: entry for entry in entries}
|
||||
self.update_calls: List[tuple] = []
|
||||
|
||||
async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool:
|
||||
self.update_calls.append((model_type, file_path, autov3))
|
||||
entry = self.entries.get(file_path)
|
||||
if entry is None:
|
||||
return False
|
||||
old_item = dict(entry)
|
||||
new_item = dict(entry)
|
||||
new_item['autov3'] = autov3
|
||||
self._persistent_cache.update_single_model(model_type, new_item, old_item)
|
||||
entry['autov3'] = autov3
|
||||
return True
|
||||
|
||||
|
||||
def _make_store(tmp_path: Path, monkeypatch, name: str = 'cache.sqlite') -> PersistentModelCache:
|
||||
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
|
||||
return PersistentModelCache(db_path=str(tmp_path / name))
|
||||
|
||||
|
||||
def _write_file(tmp_path: Path, name: str) -> str:
|
||||
path = tmp_path / name
|
||||
path.write_text(name, encoding='utf-8')
|
||||
return path.as_posix()
|
||||
|
||||
|
||||
async def test_backfill_updates_models_and_self_terminates(tmp_path: Path, monkeypatch) -> None:
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
|
||||
path_a = _write_file(tmp_path, 'a.txt')
|
||||
path_b = _write_file(tmp_path, 'b.txt')
|
||||
checked = (tmp_path / 'checked.txt').as_posix()
|
||||
valued = (tmp_path / 'valued.txt').as_posix()
|
||||
|
||||
entries = [
|
||||
_entry(path_a, 'hash-a'),
|
||||
_entry(path_b, 'hash-b'),
|
||||
_entry(checked, 'hash-checked', autov3=''),
|
||||
_entry(valued, 'hash-valued', autov3='a1b2c3d4e5f6'),
|
||||
]
|
||||
store.save_cache(
|
||||
'dummy',
|
||||
entries,
|
||||
{e['sha256']: [e['file_path']] for e in entries},
|
||||
[],
|
||||
)
|
||||
|
||||
scanner = RecordingScanner('dummy', store, entries)
|
||||
updated = await Autov3BackfillService.get_instance().backfill(scanner)
|
||||
|
||||
# Non-safetensors files yield no embedded hash, so both are marked ''.
|
||||
assert updated == 2
|
||||
assert set(scanner.update_calls) == {('dummy', path_a, ''), ('dummy', path_b, '')}
|
||||
|
||||
# Self-terminating: the driving query now finds no remaining rows.
|
||||
assert store.get_models_missing_autov3('dummy') == []
|
||||
|
||||
persisted = store.load_cache('dummy')
|
||||
items = {item['file_path']: item for item in persisted.raw_data}
|
||||
assert items[path_a]['autov3'] == ''
|
||||
assert items[path_b]['autov3'] == ''
|
||||
# Checked-unavailable and valued rows are never recomputed or touched.
|
||||
assert items[checked]['autov3'] == ''
|
||||
assert items[valued]['autov3'] == 'a1b2c3d4e5f6'
|
||||
|
||||
|
||||
async def test_backfill_skips_missing_files_without_marking(tmp_path: Path, monkeypatch) -> None:
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
|
||||
existing = _write_file(tmp_path, 'existing.txt')
|
||||
missing = (tmp_path / 'missing.txt').as_posix()
|
||||
|
||||
entries = [_entry(existing, 'hash-existing'), _entry(missing, 'hash-missing')]
|
||||
store.save_cache(
|
||||
'dummy',
|
||||
entries,
|
||||
{'hash-existing': [existing], 'hash-missing': [missing]},
|
||||
[],
|
||||
)
|
||||
|
||||
scanner = RecordingScanner('dummy', store, entries)
|
||||
updated = await Autov3BackfillService.get_instance().backfill(scanner)
|
||||
|
||||
assert updated == 1
|
||||
assert scanner.update_calls == [('dummy', existing, '')]
|
||||
# The missing row was not marked, so it still appears in the query.
|
||||
assert store.get_models_missing_autov3('dummy') == [missing]
|
||||
|
||||
|
||||
async def test_backfill_returns_zero_when_same_type_already_running(tmp_path: Path, monkeypatch) -> None:
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
scanner = RecordingScanner('dummy', store, [])
|
||||
|
||||
service = Autov3BackfillService.get_instance()
|
||||
service._running_types = {'dummy'}
|
||||
try:
|
||||
assert await service.backfill(scanner) == 0
|
||||
finally:
|
||||
service._running_types = set()
|
||||
assert scanner.update_calls == []
|
||||
|
||||
|
||||
async def test_backfill_runs_concurrently_for_different_model_types(tmp_path: Path, monkeypatch) -> None:
|
||||
"""Scanners initialize in parallel (lora_manager.py), so a backfill for one
|
||||
model type must not skip another type's backfill."""
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
lora_file = _write_file(tmp_path, 'lora.txt')
|
||||
ckpt_file = _write_file(tmp_path, 'ckpt.txt')
|
||||
store.save_cache(
|
||||
'lora',
|
||||
[_entry(lora_file, 'hash-lora')],
|
||||
{'hash-lora': [lora_file]},
|
||||
[],
|
||||
)
|
||||
store.save_cache(
|
||||
'checkpoint',
|
||||
[_entry(ckpt_file, 'hash-ckpt')],
|
||||
{'hash-ckpt': [ckpt_file]},
|
||||
[],
|
||||
)
|
||||
|
||||
lora_scanner = RecordingScanner('lora', store, [_entry(lora_file, 'hash-lora')])
|
||||
ckpt_scanner = RecordingScanner('checkpoint', store, [_entry(ckpt_file, 'hash-ckpt')])
|
||||
|
||||
service = Autov3BackfillService.get_instance()
|
||||
service._running_types = {'checkpoint'} # Simulate a checkpoint backfill in flight
|
||||
|
||||
try:
|
||||
# The lora backfill must still run while checkpoint is in progress.
|
||||
assert await service.backfill(lora_scanner) == 1
|
||||
assert lora_scanner.update_calls == [('lora', lora_file, '')]
|
||||
finally:
|
||||
service._running_types = set()
|
||||
|
||||
|
||||
async def test_backfill_never_raises_on_failure(tmp_path: Path, monkeypatch) -> None:
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
existing = _write_file(tmp_path, 'boom.txt')
|
||||
|
||||
class RaisingScanner(RecordingScanner):
|
||||
async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool:
|
||||
raise RuntimeError('boom')
|
||||
|
||||
entries = [_entry(existing, 'hash-boom')]
|
||||
store.save_cache('dummy', entries, {'hash-boom': [existing]}, [])
|
||||
|
||||
scanner = RaisingScanner('dummy', store, entries)
|
||||
updated = await Autov3BackfillService.get_instance().backfill(scanner)
|
||||
assert updated == 0
|
||||
|
||||
|
||||
async def test_backfill_uses_default_cache_when_scanner_has_none(tmp_path: Path, monkeypatch) -> None:
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
existing = _write_file(tmp_path, 'model.txt')
|
||||
entries = [_entry(existing, 'hash-x')]
|
||||
store.save_cache('dummy', entries, {'hash-x': [existing]}, [])
|
||||
|
||||
from py.services import persistent_model_cache as pmc_module
|
||||
|
||||
monkeypatch.setattr(pmc_module, 'get_persistent_cache', lambda: store)
|
||||
|
||||
class BareScanner:
|
||||
model_type = 'dummy'
|
||||
|
||||
async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool:
|
||||
entry = next(e for e in entries if e['file_path'] == file_path)
|
||||
old_item = dict(entry)
|
||||
new_item = dict(entry)
|
||||
new_item['autov3'] = autov3
|
||||
store.update_single_model(model_type, new_item, old_item)
|
||||
return True
|
||||
|
||||
updated = await Autov3BackfillService.get_instance().backfill(BareScanner())
|
||||
assert updated == 1
|
||||
assert store.get_models_missing_autov3('dummy') == []
|
||||
|
||||
|
||||
async def test_backfill_idempotent_second_run_is_noop(tmp_path: Path, monkeypatch) -> None:
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
existing = _write_file(tmp_path, 'idem.txt')
|
||||
|
||||
entries = [_entry(existing, 'hash-idem')]
|
||||
store.save_cache('dummy', entries, {'hash-idem': [existing]}, [])
|
||||
|
||||
scanner = RecordingScanner('dummy', store, entries)
|
||||
service = Autov3BackfillService.get_instance()
|
||||
|
||||
assert await service.backfill(scanner) == 1
|
||||
# A re-run has nothing left to do.
|
||||
assert await service.backfill(scanner) == 0
|
||||
assert len(scanner.update_calls) == 1
|
||||
|
||||
|
||||
async def test_backfill_end_to_end_through_scanner_lazy_import(tmp_path: Path, monkeypatch) -> None:
|
||||
"""Drive the scanner's lazy-import trigger (`_run_autov3_backfill`) end to end."""
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
|
||||
path_a = _write_file(tmp_path, 'alpha.txt')
|
||||
path_b = _write_file(tmp_path, 'beta.txt')
|
||||
|
||||
entries = [_entry(path_a, 'hash-alpha'), _entry(path_b, 'hash-beta')]
|
||||
store.save_cache(
|
||||
'dummy',
|
||||
entries,
|
||||
{'hash-alpha': [path_a], 'hash-beta': [path_b]},
|
||||
[],
|
||||
)
|
||||
|
||||
class RealScanner(ModelScanner):
|
||||
def __init__(self) -> None:
|
||||
self.model_type = 'dummy'
|
||||
self._persistent_cache = store
|
||||
self._cache = ModelCache(raw_data=[dict(e) for e in entries], folders=[])
|
||||
self._hash_index = ModelHashIndex()
|
||||
|
||||
await RealScanner()._run_autov3_backfill()
|
||||
|
||||
assert store.get_models_missing_autov3('dummy') == []
|
||||
persisted = store.load_cache('dummy')
|
||||
items = {item['file_path']: item for item in persisted.raw_data}
|
||||
assert items[path_a]['autov3'] == ''
|
||||
assert items[path_b]['autov3'] == ''
|
||||
|
||||
|
||||
async def test_backfill_prefers_civitai_autov3_from_sidecar(tmp_path: Path, monkeypatch) -> None:
|
||||
"""Backfill uses the Civitai AutoV3 for the SHA256-matching file when the
|
||||
sidecar carries Civitai metadata, even if the file itself has no embedded
|
||||
header hash (the checkpoint case)."""
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
|
||||
path = _write_file(tmp_path, 'ckpt.txt') # non-safetensors: no header hash
|
||||
sidecar = tmp_path / 'ckpt.metadata.json'
|
||||
sidecar.write_text(
|
||||
json.dumps({
|
||||
"sha256": "hash-ckpt",
|
||||
"civitai": {
|
||||
"files": [
|
||||
{"name": "other.safetensors", "hashes": {"SHA256": "zzz999"}},
|
||||
{"name": "ckpt.safetensors", "hashes": {"SHA256": "HASH-CKPT", "AutoV3": "ABCDEF1234567890"}},
|
||||
]
|
||||
},
|
||||
}),
|
||||
encoding='utf-8',
|
||||
)
|
||||
|
||||
store.save_cache('dummy', [_entry(path, 'hash-ckpt')], {'hash-ckpt': [path]}, [])
|
||||
|
||||
scanner = RecordingScanner('dummy', store, [_entry(path, 'hash-ckpt')])
|
||||
updated = await Autov3BackfillService.get_instance().backfill(scanner)
|
||||
|
||||
assert updated == 1
|
||||
assert scanner.update_calls == [('dummy', path, 'abcdef123456')]
|
||||
|
||||
persisted = store.load_cache('dummy')
|
||||
items = {item['file_path']: item for item in persisted.raw_data}
|
||||
assert items[path]['autov3'] == 'abcdef123456'
|
||||
# Self-terminating: the row is marked and the driving query empties.
|
||||
assert store.get_models_missing_autov3('dummy') == []
|
||||
|
||||
|
||||
async def test_backfill_falls_back_to_header_when_sidecar_has_no_match(tmp_path: Path, monkeypatch) -> None:
|
||||
"""When the sidecar's Civitai files do not contain a SHA256 match, the
|
||||
backfill falls back to the embedded header hash ('' for non-safetensors)."""
|
||||
store = _make_store(tmp_path, monkeypatch)
|
||||
|
||||
path = _write_file(tmp_path, 'plain.txt')
|
||||
sidecar = tmp_path / 'plain.metadata.json'
|
||||
sidecar.write_text(
|
||||
json.dumps({
|
||||
"sha256": "hash-plain",
|
||||
"civitai": {"files": [{"name": "other.safetensors", "hashes": {"SHA256": "zzz999", "AutoV3": "ABCDEF123456"}}]},
|
||||
}),
|
||||
encoding='utf-8',
|
||||
)
|
||||
|
||||
store.save_cache('dummy', [_entry(path, 'hash-plain')], {'hash-plain': [path]}, [])
|
||||
|
||||
scanner = RecordingScanner('dummy', store, [_entry(path, 'hash-plain')])
|
||||
updated = await Autov3BackfillService.get_instance().backfill(scanner)
|
||||
|
||||
assert updated == 1
|
||||
assert scanner.update_calls == [('dummy', path, '')]
|
||||
@@ -1318,3 +1318,70 @@ class TestHfGroupKey:
|
||||
"hf_url": "https://huggingface.co/user/repo",
|
||||
}
|
||||
assert BaseModelService._extract_group_key(item) == "hf:user/repo"
|
||||
|
||||
|
||||
class TestApplyHashFilters:
|
||||
"""_apply_hash_filters matches items by SHA256 or non-empty AutoV3."""
|
||||
|
||||
def _make_service(self):
|
||||
return DummyService(model_type="stub", scanner=object(), metadata_class=BaseModelMetadata)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_matches_item_by_autov3(self):
|
||||
service = self._make_service()
|
||||
data = [
|
||||
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": "abcdef123456"},
|
||||
{"file_path": "/m/two.safetensors", "sha256": "b" * 64, "autov3": ""},
|
||||
]
|
||||
|
||||
result = await service._apply_hash_filters(data, {"single_hash": "ABCDEF123456"})
|
||||
|
||||
assert [item["file_path"] for item in result] == ["/m/one.safetensors"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_matches_item_by_sha256(self):
|
||||
service = self._make_service()
|
||||
data = [
|
||||
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": ""},
|
||||
]
|
||||
|
||||
result = await service._apply_hash_filters(data, {"single_hash": "A" * 64})
|
||||
|
||||
assert [item["file_path"] for item in result] == ["/m/one.safetensors"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_or_absent_autov3_never_matches(self):
|
||||
service = self._make_service()
|
||||
data = [
|
||||
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": ""},
|
||||
{"file_path": "/m/two.safetensors", "sha256": "b" * 64},
|
||||
]
|
||||
|
||||
result = await service._apply_hash_filters(data, {"single_hash": "cdef123456ab"})
|
||||
|
||||
assert result == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_hashes_match_autov3_and_sha256(self):
|
||||
service = self._make_service()
|
||||
data = [
|
||||
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": "abcdef123456"},
|
||||
{"file_path": "/m/two.safetensors", "sha256": "b" * 64, "autov3": ""},
|
||||
]
|
||||
|
||||
result = await service._apply_hash_filters(
|
||||
data, {"multiple_hashes": ["abcdef123456", "c" * 64]}
|
||||
)
|
||||
|
||||
assert [item["file_path"] for item in result] == ["/m/one.safetensors"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_hash_filters_returns_data_unchanged(self):
|
||||
service = self._make_service()
|
||||
data = [
|
||||
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": "abcdef123456"},
|
||||
]
|
||||
|
||||
result = await service._apply_hash_filters(data, {})
|
||||
|
||||
assert result == data
|
||||
|
||||
@@ -321,3 +321,91 @@ class TestCacheEntryValidator:
|
||||
|
||||
assert result.is_valid is True
|
||||
assert result.repaired is False
|
||||
|
||||
|
||||
class TestAutov3Validation:
|
||||
"""AutoV3 optional-field validation semantics."""
|
||||
|
||||
def _entry(self, **overrides):
|
||||
# Fully-populated entry so that autov3 is the only candidate repair.
|
||||
entry = {
|
||||
'file_path': '/models/test.safetensors',
|
||||
'sha256': 'abc123',
|
||||
'file_name': 'test.safetensors',
|
||||
'model_name': 'Test Model',
|
||||
'folder': 'test_folder',
|
||||
'size': 1024,
|
||||
'modified': 1234567890.0,
|
||||
'tags': ['tag1'],
|
||||
'preview_url': 'http://example.com/preview.jpg',
|
||||
'base_model': 'SD1.5',
|
||||
'from_civitai': True,
|
||||
'favorite': True,
|
||||
'exclude': False,
|
||||
'db_checked': True,
|
||||
'preview_nsfw_level': 1,
|
||||
'notes': 'Test notes',
|
||||
'usage_tips': 'Test tips',
|
||||
'hash_status': 'completed',
|
||||
}
|
||||
entry.update(overrides)
|
||||
return entry
|
||||
|
||||
def test_validate_valid_autov3_normalized_to_lowercase(self):
|
||||
"""Uppercase 12-hex autov3 is normalized to lowercase under auto_repair."""
|
||||
result = CacheEntryValidator.validate(
|
||||
self._entry(autov3='ABCDEF123456'), auto_repair=True
|
||||
)
|
||||
|
||||
assert result.is_valid is True
|
||||
assert result.entry['autov3'] == 'abcdef123456'
|
||||
assert result.repaired is True
|
||||
|
||||
def test_validate_autov3_empty_string_is_valid(self):
|
||||
"""Empty autov3 means checked-but-unavailable and is valid."""
|
||||
result = CacheEntryValidator.validate(
|
||||
self._entry(autov3=''), auto_repair=False
|
||||
)
|
||||
|
||||
assert result.is_valid is True
|
||||
assert result.repaired is False
|
||||
|
||||
def test_validate_autov3_none_is_valid_and_not_counted_as_repair(self):
|
||||
"""autov3 None (not checked) is valid and is NOT counted as a repair."""
|
||||
result = CacheEntryValidator.validate(
|
||||
self._entry(autov3=None), auto_repair=True
|
||||
)
|
||||
|
||||
assert result.is_valid is True
|
||||
assert result.repaired is False
|
||||
assert result.entry['autov3'] is None
|
||||
|
||||
def test_validate_absent_autov3_is_valid_and_not_counted_as_repair(self):
|
||||
"""A missing autov3 field is valid and is NOT counted as a repair."""
|
||||
result = CacheEntryValidator.validate(self._entry(), auto_repair=True)
|
||||
|
||||
assert result.is_valid is True
|
||||
assert result.repaired is False
|
||||
assert 'autov3' not in result.entry
|
||||
|
||||
def test_validate_short_autov3_still_valid_and_repaired_to_none(self):
|
||||
"""A malformed autov3 does not invalidate the entry (optional field);
|
||||
with auto_repair the value is repaired to None."""
|
||||
result = CacheEntryValidator.validate(
|
||||
self._entry(autov3='abc'), auto_repair=True
|
||||
)
|
||||
|
||||
assert result.is_valid is True
|
||||
assert result.entry['autov3'] is None
|
||||
assert result.repaired is True
|
||||
|
||||
def test_validate_non_string_autov3_still_valid_and_repaired_to_none(self):
|
||||
"""A non-string autov3 does not invalidate the entry (optional field);
|
||||
with auto_repair the value is repaired to None."""
|
||||
result = CacheEntryValidator.validate(
|
||||
self._entry(autov3=123), auto_repair=True
|
||||
)
|
||||
|
||||
assert result.is_valid is True
|
||||
assert result.entry['autov3'] is None
|
||||
assert result.repaired is True
|
||||
|
||||
@@ -112,6 +112,68 @@ async def test_update_model_metadata_merges_and_persists():
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_model_metadata_propagates_civitai_autov3():
|
||||
helpers = build_service()
|
||||
|
||||
local = {
|
||||
"sha256": "111aabbf94dd9e59c05d842fccf57bec915b2a3c237f6b54f8d614e40858d717",
|
||||
"autov3": "",
|
||||
"model_name": "Local",
|
||||
}
|
||||
remote = {
|
||||
"source": "api",
|
||||
"model": {"name": "Remote Model", "description": "", "tags": []},
|
||||
"images": [],
|
||||
"files": [
|
||||
{
|
||||
"name": "other.safetensors",
|
||||
"hashes": {"SHA256": "ZZZ999"},
|
||||
},
|
||||
{
|
||||
"name": "model.safetensors",
|
||||
"hashes": {
|
||||
"SHA256": "111aabbf94dd9e59c05d842fccf57bec915b2a3c237f6b54f8d614e40858d717",
|
||||
"AutoV3": "8A582E901D7F",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
result = await helpers.service.update_model_metadata(
|
||||
"path/to/model.metadata.json",
|
||||
local,
|
||||
remote,
|
||||
helpers.default_provider,
|
||||
)
|
||||
|
||||
# Civitai-first: the '' (checked-unavailable) state is upgraded in-session
|
||||
# by the freshly fetched metadata, without any header re-read.
|
||||
assert result["autov3"] == "8a582e901d7f"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_model_metadata_keeps_autov3_without_matching_file():
|
||||
helpers = build_service()
|
||||
|
||||
local = {"sha256": "abc123", "autov3": "", "model_name": "Local"}
|
||||
remote = {
|
||||
"source": "api",
|
||||
"model": {"name": "Remote Model", "description": "", "tags": []},
|
||||
"images": [],
|
||||
"files": [{"name": "other.safetensors", "hashes": {"SHA256": "ZZZ999", "AutoV3": "ABCDEF123456"}}],
|
||||
}
|
||||
|
||||
result = await helpers.service.update_model_metadata(
|
||||
"path/to/model.metadata.json",
|
||||
local,
|
||||
remote,
|
||||
helpers.default_provider,
|
||||
)
|
||||
|
||||
assert result["autov3"] == ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_and_update_model_success_updates_cache(tmp_path):
|
||||
helpers = build_service()
|
||||
|
||||
@@ -111,3 +111,145 @@ class TestModelHashIndexGetDuplicateFilenames:
|
||||
index.add_entry("abc123", "/a/lora.safetensors")
|
||||
assert len(index) == 1
|
||||
assert index.get_duplicate_filenames() == {}
|
||||
|
||||
|
||||
class TestModelHashIndexAutov3:
|
||||
"""AutoV3 hash index behavior."""
|
||||
|
||||
def test_add_entry_with_autov3_supports_lookup_by_autov3(self):
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="AbCdEf123456")
|
||||
|
||||
assert index.has_hash("abcdef123456") is True
|
||||
assert index.get_path("abcdef123456") == "/models/lora.safetensors"
|
||||
assert index.get_all_autov3() == {"abcdef123456": "/models/lora.safetensors"}
|
||||
|
||||
def test_add_entry_without_autov3_creates_no_autov3_lookup(self):
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("b" * 64, "/models/lora.safetensors")
|
||||
|
||||
assert index.has_hash("abcdef123456") is False
|
||||
assert index.get_path("abcdef123456") is None
|
||||
assert index.get_all_autov3() == {}
|
||||
|
||||
def test_add_autov3_standalone_supports_lookup(self):
|
||||
index = ModelHashIndex()
|
||||
index.add_autov3("cdef123456ab", "/models/only_autov3.safetensors")
|
||||
|
||||
assert index.has_hash("cdef123456ab") is True
|
||||
assert index.get_path("cdef123456ab") == "/models/only_autov3.safetensors"
|
||||
assert index.get_all_autov3() == {"cdef123456ab": "/models/only_autov3.safetensors"}
|
||||
|
||||
def test_remove_by_path_removes_autov3_mapping(self):
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
|
||||
|
||||
index.remove_by_path("/models/lora.safetensors")
|
||||
|
||||
assert index.has_hash("abcdef123456") is False
|
||||
assert index.get_all_autov3() == {}
|
||||
|
||||
def test_remove_by_hash_removes_autov3_mapping(self):
|
||||
index = ModelHashIndex()
|
||||
sha256 = "a" * 64
|
||||
index.add_entry(sha256, "/models/lora.safetensors", autov3="abcdef123456")
|
||||
|
||||
index.remove_by_hash(sha256)
|
||||
|
||||
assert index.has_hash("abcdef123456") is False
|
||||
assert index.get_all_autov3() == {}
|
||||
|
||||
def test_clear_empties_autov3_index(self):
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("a" * 64, "/models/a.safetensors", autov3="aaaaabbbbbcc")
|
||||
index.add_entry("b" * 64, "/models/b.safetensors", autov3="dddddeeeeeff")
|
||||
|
||||
index.clear()
|
||||
|
||||
assert index.get_all_autov3() == {}
|
||||
assert index.has_hash("aaaaabbbbbcc") is False
|
||||
|
||||
def test_same_autov3_last_write_wins(self):
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("a" * 64, "/models/first.safetensors", autov3="abcdef123456")
|
||||
index.add_entry("b" * 64, "/models/second.safetensors", autov3="abcdef123456")
|
||||
|
||||
assert index.get_path("abcdef123456") == "/models/second.safetensors"
|
||||
assert index.get_all_autov3() == {"abcdef123456": "/models/second.safetensors"}
|
||||
|
||||
def test_dispatch_len_10_hits_autov2(self):
|
||||
index = ModelHashIndex()
|
||||
sha256 = "a" * 64
|
||||
index.add_entry(sha256, "/models/lora.safetensors")
|
||||
|
||||
assert index.get_path(sha256[:10]) == "/models/lora.safetensors"
|
||||
assert index.has_hash(sha256[:10]) is True
|
||||
|
||||
def test_dispatch_len_64_hits_sha256(self):
|
||||
index = ModelHashIndex()
|
||||
sha256 = "b" * 64
|
||||
index.add_entry(sha256, "/models/lora.safetensors")
|
||||
|
||||
assert index.get_path(sha256) == "/models/lora.safetensors"
|
||||
assert index.has_hash(sha256) is True
|
||||
|
||||
def test_dispatch_len_12_hits_autov3(self):
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("c" * 64, "/models/lora.safetensors", autov3="cdef123456ab")
|
||||
|
||||
assert index.get_path("cdef123456ab") == "/models/lora.safetensors"
|
||||
assert index.has_hash("cdef123456ab") is True
|
||||
|
||||
def test_add_entry_drops_stale_autov3_for_replaced_path(self):
|
||||
# A file replaced in place (new content → new sha256 and new autov3)
|
||||
# must not keep the old autov3 mapping — it would survive into the
|
||||
# persisted snapshot and make lookups resolve the wrong file.
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
|
||||
index.add_entry("b" * 64, "/models/lora.safetensors", autov3="fedcba654321")
|
||||
|
||||
assert index.get_path("abcdef123456") is None
|
||||
assert index.has_hash("abcdef123456") is False
|
||||
assert index.get_path("fedcba654321") == "/models/lora.safetensors"
|
||||
assert index.get_all_autov3() == {"fedcba654321": "/models/lora.safetensors"}
|
||||
|
||||
def test_add_entry_without_autov3_drops_stale_mapping_for_replaced_path(self):
|
||||
# Replaced file whose new content has no embedded hash: the stale
|
||||
# autov3 mapping must be dropped, not left pointing at the path.
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
|
||||
index.add_entry("b" * 64, "/models/lora.safetensors")
|
||||
|
||||
assert index.get_path("abcdef123456") is None
|
||||
assert index.get_all_autov3() == {}
|
||||
|
||||
def test_add_entry_re_registration_with_same_autov3_is_idempotent(self):
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
|
||||
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
|
||||
|
||||
assert index.get_path("abcdef123456") == "/models/lora.safetensors"
|
||||
assert index.get_all_autov3() == {"abcdef123456": "/models/lora.safetensors"}
|
||||
|
||||
def test_add_entry_same_sha_without_autov3_preserves_existing_mapping(self):
|
||||
# A lazy-hash completion (checkpoint_scanner) re-registers the SAME
|
||||
# file with the same sha256 but omits autov3. That must never clear
|
||||
# the previously registered autov3 mapping.
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("a" * 64, "/models/ckpt.safetensors", autov3="abcdef123456")
|
||||
index.add_entry("a" * 64, "/models/ckpt.safetensors")
|
||||
|
||||
assert index.get_path("abcdef123456") == "/models/ckpt.safetensors"
|
||||
assert index.get_all_autov3() == {"abcdef123456": "/models/ckpt.safetensors"}
|
||||
|
||||
def test_add_entry_same_sha_with_new_autov3_drops_old_mapping(self):
|
||||
# Re-registration with an explicit, different autov3 (metadata
|
||||
# correction) must drop the stale mapping for that path.
|
||||
index = ModelHashIndex()
|
||||
index.add_entry("a" * 64, "/models/ckpt.safetensors", autov3="abcdef123456")
|
||||
index.add_entry("a" * 64, "/models/ckpt.safetensors", autov3="fedcba654321")
|
||||
|
||||
assert index.get_path("abcdef123456") is None
|
||||
assert index.has_hash("abcdef123456") is False
|
||||
assert index.get_path("fedcba654321") == "/models/ckpt.safetensors"
|
||||
assert index.get_all_autov3() == {"fedcba654321": "/models/ckpt.safetensors"}
|
||||
|
||||
@@ -341,3 +341,112 @@ def test_update_single_model_update_hash(tmp_path: Path, monkeypatch):
|
||||
new_hash_pairs = [p for p in persisted.hash_rows if p[0] == 'new-hash']
|
||||
assert len(new_hash_pairs) == 1
|
||||
assert new_hash_pairs[0][1] == file_path
|
||||
|
||||
|
||||
# ── get_models_missing_autov3 ─────────────────────────────────────────
|
||||
|
||||
|
||||
def _autov3_entry(file_path: str, sha256: str, autov3=None) -> dict:
|
||||
"""Minimal model entry for the models table (autov3 tri-state preserved)."""
|
||||
return {
|
||||
'file_path': file_path,
|
||||
'file_name': Path(file_path).stem,
|
||||
'model_name': Path(file_path).stem,
|
||||
'folder': '',
|
||||
'size': 1,
|
||||
'modified': 1.0,
|
||||
'sha256': sha256,
|
||||
'autov3': autov3,
|
||||
'base_model': '',
|
||||
'preview_url': '',
|
||||
'preview_nsfw_level': 0,
|
||||
'from_civitai': True,
|
||||
'favorite': False,
|
||||
'notes': '',
|
||||
'usage_tips': '',
|
||||
'metadata_source': None,
|
||||
'exclude': False,
|
||||
'db_checked': False,
|
||||
'last_checked_at': 0.0,
|
||||
'tags': [],
|
||||
'civitai': None,
|
||||
'civitai_deleted': False,
|
||||
'skip_metadata_refresh': False,
|
||||
'license_flags': DEFAULT_LICENSE_FLAGS,
|
||||
'hash_status': 'completed',
|
||||
'hf_url': '',
|
||||
}
|
||||
|
||||
|
||||
def test_get_models_missing_autov3_filters_rows(tmp_path: Path, monkeypatch) -> None:
|
||||
"""Only NULL-autov3 rows with a completed sha256 qualify."""
|
||||
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
|
||||
db_path = tmp_path / 'cache.sqlite'
|
||||
store = PersistentModelCache(db_path=str(db_path))
|
||||
|
||||
null_path = (tmp_path / 'null.txt').as_posix()
|
||||
checked_path = (tmp_path / 'checked.txt').as_posix()
|
||||
valued_path = (tmp_path / 'valued.txt').as_posix()
|
||||
empty_sha_path = (tmp_path / 'empty_sha.txt').as_posix()
|
||||
|
||||
store.save_cache(
|
||||
'dummy',
|
||||
[
|
||||
_autov3_entry(null_path, 'hash-null'),
|
||||
_autov3_entry(checked_path, 'hash-checked', autov3=''),
|
||||
_autov3_entry(valued_path, 'hash-valued', autov3='a1b2c3d4e5f6'),
|
||||
_autov3_entry(empty_sha_path, ''),
|
||||
],
|
||||
{},
|
||||
[],
|
||||
)
|
||||
|
||||
assert store.get_models_missing_autov3('dummy') == [null_path]
|
||||
|
||||
|
||||
def test_get_models_missing_autov3_filters_by_model_type(tmp_path: Path, monkeypatch) -> None:
|
||||
"""Only rows of the requested model_type are returned."""
|
||||
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
|
||||
db_path = tmp_path / 'cache.sqlite'
|
||||
store = PersistentModelCache(db_path=str(db_path))
|
||||
|
||||
lora_path = (tmp_path / 'lora.txt').as_posix()
|
||||
checkpoint_path = (tmp_path / 'checkpoint.txt').as_posix()
|
||||
|
||||
store.save_cache('lora', [_autov3_entry(lora_path, 'hash-lora')], {}, [])
|
||||
store.save_cache('checkpoint', [_autov3_entry(checkpoint_path, 'hash-checkpoint')], {}, [])
|
||||
|
||||
assert store.get_models_missing_autov3('lora') == [lora_path]
|
||||
assert store.get_models_missing_autov3('checkpoint') == [checkpoint_path]
|
||||
|
||||
|
||||
def test_get_models_missing_autov3_empty_on_clean_db(tmp_path: Path, monkeypatch) -> None:
|
||||
"""A freshly created database has no rows to backfill."""
|
||||
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
|
||||
store = PersistentModelCache(db_path=str(tmp_path / 'cache.sqlite'))
|
||||
assert store.get_models_missing_autov3('dummy') == []
|
||||
|
||||
|
||||
def test_get_models_missing_autov3_disabled_cache_returns_empty(tmp_path: Path, monkeypatch) -> None:
|
||||
"""When the persistent cache is disabled the query is a no-op."""
|
||||
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '1')
|
||||
store = PersistentModelCache(db_path=str(tmp_path / 'cache.sqlite'))
|
||||
assert store.get_models_missing_autov3('dummy') == []
|
||||
|
||||
|
||||
def test_get_models_missing_autov3_self_terminates_after_marking(tmp_path: Path, monkeypatch) -> None:
|
||||
"""Once a row receives a checked state it drops out of the query."""
|
||||
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
|
||||
db_path = tmp_path / 'cache.sqlite'
|
||||
store = PersistentModelCache(db_path=str(db_path))
|
||||
|
||||
file_path = (tmp_path / 'm.txt').as_posix()
|
||||
store.save_cache('dummy', [_autov3_entry(file_path, 'hash-m')], {}, [])
|
||||
assert store.get_models_missing_autov3('dummy') == [file_path]
|
||||
|
||||
# Mark the row '' (checked-unavailable) and re-query.
|
||||
old_item = {'file_path': file_path, 'tags': [], 'sha256': 'hash-m'}
|
||||
new_item = _autov3_entry(file_path, 'hash-m', autov3='')
|
||||
store.update_single_model('dummy', new_item, old_item=old_item)
|
||||
|
||||
assert store.get_models_missing_autov3('dummy') == []
|
||||
|
||||
Reference in New Issue
Block a user