ComfyUI-Lora-Manager/tests/services/test_persistent_model_cache.py

from pathlib import Path
from typing import Any, Dict

import pytest

from py.services.persistent_model_cache import PersistentModelCache, DEFAULT_LICENSE_FLAGS


def test_persistent_cache_roundtrip(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
    db_path = tmp_path / 'cache.sqlite'
    store = PersistentModelCache(db_path=str(db_path))

    file_a = (tmp_path / 'a.txt').as_posix()
    file_b = (tmp_path / 'b.txt').as_posix()
    duplicate_path = f"{file_b}.copy"

    raw_data = [
        {
            'file_path': file_a,
            'file_name': 'a',
            'model_name': 'Model A',
            'folder': '',
            'size': 10,
            'modified': 100.0,
            'sha256': 'hash-a',
            'base_model': 'base',
            'preview_url': 'preview/a.png',
            'preview_nsfw_level': 1,
            'from_civitai': True,
            'favorite': True,
            'notes': 'note',
            'usage_tips': '{}',
            'metadata_source': 'civitai_api',
            'exclude': False,
            'db_checked': True,
            'last_checked_at': 200.0,
            'tags': ['alpha', 'beta'],
            'civitai_deleted': False,
            'civitai': {
                'id': 1,
                'modelId': 2,
                'name': 'verA',
                'trainedWords': ['word1'],
                'creator': {'username': 'artist42'},
            },
            'license_flags': 13,
        },
        {
            'file_path': file_b,
            'file_name': 'b',
            'model_name': 'Model B',
            'folder': 'folder',
            'size': 20,
            'modified': 120.0,
            'sha256': 'hash-b',
            'base_model': '',
            'preview_url': '',
            'preview_nsfw_level': 0,
            'from_civitai': False,
            'favorite': False,
            'notes': '',
            'usage_tips': '',
            'metadata_source': None,
            'exclude': True,
            'db_checked': False,
            'last_checked_at': 0.0,
            'tags': [],
            'civitai_deleted': True,
            'civitai': None,
        },
    ]

    hash_index = {
        'hash-a': [file_a],
        'hash-b': [file_b, duplicate_path],
    }
    excluded = [duplicate_path]

    store.save_cache('dummy', raw_data, hash_index, excluded)

    persisted = store.load_cache('dummy')
    assert persisted is not None
    assert len(persisted.raw_data) == 2

    items = {item['file_path']: item for item in persisted.raw_data}
    assert set(items.keys()) == {file_a, file_b}
    first = items[file_a]
    assert first['favorite'] is True
    assert first['civitai']['id'] == 1
    assert first['civitai']['trainedWords'] == ['word1']
    assert first['tags'] == ['alpha', 'beta']
    assert first['metadata_source'] == 'civitai_api'
    assert first['civitai']['creator']['username'] == 'artist42'
    assert first['civitai_deleted'] is False
    assert first['license_flags'] == 13

    second = items[file_b]
    assert second['exclude'] is True
    assert second['civitai'] is None
    assert second['metadata_source'] is None
    assert second['civitai_deleted'] is True
    assert second['license_flags'] == DEFAULT_LICENSE_FLAGS

    expected_hash_pairs = {
        ('hash-a', file_a),
        ('hash-b', file_b),
        ('hash-b', duplicate_path),
    }
    assert set((sha, path) for sha, path in persisted.hash_rows) == expected_hash_pairs
    assert persisted.excluded_models == excluded


def test_incremental_updates_only_touch_changed_rows(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
    db_path = tmp_path / 'cache.sqlite'
    store = PersistentModelCache(db_path=str(db_path))

    file_a = (tmp_path / 'a.txt').as_posix()
    file_b = (tmp_path / 'b.txt').as_posix()

    initial_payload = [
        {
            'file_path': file_a,
            'file_name': 'a',
            'model_name': 'Model A',
            'folder': '',
            'size': 10,
            'modified': 100.0,
            'sha256': 'hash-a',
            'base_model': 'base',
            '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': ['alpha'],
            'civitai_deleted': False,
            'civitai': None,
        },
        {
            'file_path': file_b,
            'file_name': 'b',
            'model_name': 'Model B',
            'folder': '',
            'size': 20,
            'modified': 120.0,
            'sha256': 'hash-b',
            'base_model': '',
            'preview_url': '',
            'preview_nsfw_level': 0,
            'from_civitai': False,
            'favorite': False,
            'notes': '',
            'usage_tips': '',
            'metadata_source': 'civarchive',
            'exclude': False,
            'db_checked': False,
            'last_checked_at': 0.0,
            'tags': ['beta'],
            'civitai_deleted': False,
            'civitai': {'creator': {'username': 'builder'}},
        },
    ]

    statements: list[str] = []
    original_connect = store._connect

    def _recording_connect(readonly: bool = False):
        conn = original_connect(readonly=readonly)
        conn.set_trace_callback(statements.append)
        return conn

    store._connect = _recording_connect  # type: ignore[method-assign]

    store.save_cache('dummy', initial_payload, {'hash-a': [file_a], 'hash-b': [file_b]}, [])
    statements.clear()

    updated_payload = [
        initial_payload[0],
        {
            **initial_payload[1],
            'model_name': 'Model B Updated',
            'favorite': True,
            'tags': ['beta', 'gamma'],
            'metadata_source': 'archive_db',
            'civitai_deleted': True,
            'civitai': {'creator': {'username': 'builder_v2'}},
        },
    ]
    hash_index = {'hash-a': [file_a], 'hash-b': [file_b]}

    store.save_cache('dummy', updated_payload, hash_index, [])

    broad_delete = [
        stmt for stmt in statements if "DELETE FROM models WHERE model_type = 'dummy'" in stmt and "file_path" not in stmt
    ]
    assert not broad_delete

    updated_stmt_present = any(
        "UPDATE models" in stmt and f"file_path = '{file_b}'" in stmt for stmt in statements
    )
    assert updated_stmt_present

    unchanged_stmt_present = any(
        "UPDATE models" in stmt and f"file_path = '{file_a}'" in stmt for stmt in statements
    )
    assert not unchanged_stmt_present

    tag_insert = any(
        "INSERT INTO model_tags" in stmt and "gamma" in stmt for stmt in statements
    )
    assert tag_insert

    assert any("metadata_source" in stmt for stmt in statements if "UPDATE models" in stmt)

    persisted = store.load_cache('dummy')
    assert persisted is not None
    items = {item['file_path']: item for item in persisted.raw_data}
    second = items[file_b]
    assert second['metadata_source'] == 'archive_db'
    assert second['civitai_deleted'] is True
    assert second['civitai']['creator']['username'] == 'builder_v2'


# ── update_single_model ───────────────────────────────────────────────


def test_update_single_model_insert(tmp_path: Path, monkeypatch):
    """Insert a brand-new model row via update_single_model."""
    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 / 'x.safetensors').as_posix()
    new_item = {
        'file_path': file_path,
        'file_name': 'x',
        'model_name': 'Model X',
        'folder': '',
        'size': 42,
        'modified': 1.0,
        'sha256': 'sha-x',
        'base_model': 'SDXL',
        'preview_url': '',
        'preview_nsfw_level': 0,
        'from_civitai': True,
        'favorite': True,
        'notes': 'test note',
        'usage_tips': '{}',
        'metadata_source': None,
        'exclude': False,
        'db_checked': False,
        'last_checked_at': 0.0,
        'tags': ['test', 'new'],
        'civitai': None,
        'civitai_deleted': False,
        'skip_metadata_refresh': False,
        'license_flags': DEFAULT_LICENSE_FLAGS,
        'hash_status': 'completed',
        'hf_url': '',
    }

    store.update_single_model('dummy', new_item)

    persisted = store.load_cache('dummy')
    assert persisted is not None
    items = {item['file_path']: item for item in persisted.raw_data}
    assert file_path in items
    assert items[file_path]['model_name'] == 'Model X'
    assert items[file_path]['favorite'] is True
    assert sorted(items[file_path]['tags']) == ['new', 'test']


def test_update_single_model_update_tags(tmp_path: Path, monkeypatch):
    """Tags are updated incrementally: old tags removed, new tags added."""
    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 / 'y.safetensors').as_posix()
    base = {
        'file_path': file_path, 'file_name': 'y', 'model_name': 'Y',
        'folder': '', 'size': 1, 'modified': 1.0, 'sha256': 'sha-y',
        '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, 'civitai': None, 'civitai_deleted': False,
        'skip_metadata_refresh': False, 'license_flags': DEFAULT_LICENSE_FLAGS,
        'hash_status': 'completed', 'hf_url': '',
    }

    # First insert with tags [alpha, beta]
    store.update_single_model('dummy', {**base, 'tags': ['alpha', 'beta']})

    # Now update: replace with [beta, gamma]
    old_item = {'file_path': file_path, 'tags': ['alpha', 'beta'], 'sha256': 'sha-y'}
    new_item = {**base, 'tags': ['beta', 'gamma']}
    store.update_single_model('dummy', new_item, old_item=old_item)

    persisted = store.load_cache('dummy')
    assert persisted is not None
    items = {item['file_path']: item for item in persisted.raw_data}
    assert sorted(items[file_path]['tags']) == ['beta', 'gamma']


def test_update_single_model_update_hash(tmp_path: Path, monkeypatch):
    """When sha256 changes, the hash_index is updated incrementally."""
    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 / 'z.safetensors').as_posix()
    base = {
        'file_path': file_path, 'file_name': 'z', 'model_name': 'Z',
        'folder': '', 'size': 1, 'modified': 1.0, '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': '',
    }

    store.update_single_model('dummy', {**base, 'sha256': 'old-hash'})

    old_item = {'file_path': file_path, 'tags': [], 'sha256': 'old-hash'}
    new_item = {**base, 'sha256': 'new-hash'}
    store.update_single_model('dummy', new_item, old_item=old_item)

    persisted = store.load_cache('dummy')
    assert persisted is not None
    # old hash should be gone from hash_index
    old_hash_pairs = [p for p in persisted.hash_rows if p[0] == 'old-hash']
    assert len(old_hash_pairs) == 0
    # new hash should be present
    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[str, Any]:
    """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') == []


# --- which roots still hold cached entries (issue #1146) --------------------


def _cache_with_paths(tmp_path: Path, monkeypatch, entries) -> PersistentModelCache:
    """Build a cache holding ``entries`` (``(model_type, file_path)`` pairs)."""

    import sqlite3

    monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
    db_path = tmp_path / 'roots.sqlite'
    store = PersistentModelCache(db_path=str(db_path))
    with sqlite3.connect(str(db_path)) as conn:
        conn.executemany(
            "INSERT INTO models (model_type, file_path) VALUES (?, ?)", entries
        )
        conn.commit()
    return store


def test_roots_with_cached_entries_matches_the_hosts_slash_style(
    tmp_path: Path, monkeypatch
) -> None:
    """The configured root may be spelled with backslashes (Windows)."""

    store = _cache_with_paths(
        tmp_path,
        monkeypatch,
        [("lora", "E:/StableDiffusion/cc_models/loras/pack/a.safetensors")],
    )

    matched = store.roots_with_cached_entries(
        "lora", ["E:\\StableDiffusion\\cc_models\\loras", "E:\\empty\\loras"]
    )

    assert matched == {"E:\\StableDiffusion\\cc_models\\loras"}


def test_roots_with_cached_entries_compares_whole_path_segments(
    tmp_path: Path, monkeypatch
) -> None:
    """``loras-extra`` and ``loras2`` are not inside ``loras``."""

    store = _cache_with_paths(
        tmp_path,
        monkeypatch,
        [
            ("lora", "E:/loras-extra/a.safetensors"),
            ("lora", "E:/loras2/b.safetensors"),
            ("checkpoint", "E:/loras/c.safetensors"),
        ],
    )

    assert store.roots_with_cached_entries("lora", ["E:/loras"]) == set()
    assert store.roots_with_cached_entries("checkpoint", ["E:/loras"]) == {"E:/loras"}


def test_roots_with_cached_entries_escapes_like_wildcards(
    tmp_path: Path, monkeypatch
) -> None:
    """A directory name may contain ``_``/``%`` without becoming a wildcard."""

    store = _cache_with_paths(
        tmp_path,
        monkeypatch,
        [("lora", "E:/myXloras/a.safetensors"), ("lora", "E:/100%/a.safetensors")],
    )

    assert store.roots_with_cached_entries("lora", ["E:/my_loras"]) == set()
    assert store.roots_with_cached_entries("lora", ["E:/100x"]) == set()
    assert store.roots_with_cached_entries("lora", ["E:/100%"]) == {"E:/100%"}


def test_roots_with_cached_entries_raises_when_unknown(
    tmp_path: Path, monkeypatch
) -> None:
    """No cache, no answer: the caller must not drop roots on a guess."""

    monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '1')
    disabled = PersistentModelCache(db_path=str(tmp_path / 'disabled.sqlite'))
    with pytest.raises(RuntimeError):
        disabled.roots_with_cached_entries("lora", ["E:/loras"])

    monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
    missing = PersistentModelCache(db_path=str(tmp_path / 'missing.sqlite'))
    (tmp_path / 'missing.sqlite').unlink()
    with pytest.raises(FileNotFoundError):
        missing.roots_with_cached_entries("lora", ["E:/loras"])


def test_roots_with_cached_entries_without_candidates(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
    store = PersistentModelCache(db_path=str(tmp_path / 'empty.sqlite'))

    assert store.roots_with_cached_entries("lora", []) == set()
    assert store.roots_with_cached_entries("lora", ["", "   "]) == set()