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

"""Centralized sidecar storage: lifecycle flows (delete/move/rename/scans)."""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any, Dict, List
from unittest.mock import AsyncMock

import pytest

from py.config import config
from py.services.checkpoint_scanner import CheckpointScanner
from py.services.model_lifecycle_service import (
    ModelLifecycleService,
    delete_model_artifacts,
)
from py.services.pending_delete_service import PendingDeleteService
from py.services.settings_manager import get_settings_manager
from py.utils.metadata_manager import MetadataManager
from py.utils.sidecar_paths import reset_root_map_cache, root_mirror_component


def _normalize(path) -> str:
    return str(path).replace(os.sep, "/")


@pytest.fixture(autouse=True)
def _isolate_root_map():
    """Root identities are cached per sidecar root; keep tests independent."""

    reset_root_map_cache()
    yield
    reset_root_map_cache()


@pytest.fixture
def library_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
    """Single checkpoint root under tmp_path; every other root emptied."""

    root = tmp_path / "checkpoints"
    root.mkdir()
    for attr, value in (
        ("loras_roots", []),
        ("base_models_roots", [str(root)]),
        ("checkpoints_roots", [str(root)]),
        ("embeddings_roots", []),
        ("other_roots", []),
        ("extra_loras_roots", []),
        ("extra_checkpoints_roots", []),
        ("extra_unet_roots", []),
        ("extra_embeddings_roots", []),
    ):
        monkeypatch.setattr(config, attr, value, raising=False)
    return root


@pytest.fixture
def centralized(library_root: Path, tmp_path: Path) -> Path:
    """Enable centralized mode rooted at tmp_path/sidecars."""

    sidecar_root = tmp_path / "sidecars"
    settings = get_settings_manager()
    settings.set("sidecar_storage_mode", "centralized")
    settings.set("sidecar_storage_path", str(sidecar_root))
    return sidecar_root


def _mirror_dir(library_root: Path, sidecar_root: Path, *rel: str) -> Path:
    """Expected mirror directory for a root-relative path."""

    component = root_mirror_component(str(library_root), sidecar_root=str(sidecar_root))
    return sidecar_root.joinpath(component, *rel)


def _write_sidecar(
    mirror_dir: Path,
    stem: str,
    *,
    file_path: str,
    preview_name: str | None = None,
    extra: Dict[str, Any] | None = None,
) -> Path:
    mirror_dir.mkdir(parents=True, exist_ok=True)
    payload: Dict[str, Any] = {
        "file_name": stem,
        "file_path": _normalize(file_path),
    }
    if preview_name:
        payload["preview_url"] = _normalize(mirror_dir / preview_name)
        (mirror_dir / preview_name).write_bytes(b"preview")
    if extra:
        payload.update(extra)
    metadata_path = mirror_dir / f"{stem}.metadata.json"
    metadata_path.write_text(json.dumps(payload), encoding="utf-8")
    return metadata_path


@pytest.mark.asyncio
async def test_delete_model_artifacts_centralized(
    library_root: Path, centralized: Path
):
    model = library_root / "model.safetensors"
    model.write_bytes(b"weights")
    mirror = _mirror_dir(library_root, centralized)
    _write_sidecar(mirror, "model", file_path=model, preview_name="model.preview.webp")

    deleted = await delete_model_artifacts(str(library_root), "model")

    assert not model.exists()
    assert not (mirror / "model.metadata.json").exists()
    assert not (mirror / "model.preview.webp").exists()
    assert any(path.endswith("model.safetensors") for path in deleted)
    assert "model.metadata.json" in deleted
    assert "model.preview.webp" in deleted


@pytest.mark.asyncio
async def test_delete_model_artifacts_alongside_still_siblings(
    library_root: Path,
):
    """Alongside mode (default) keeps deleting sidecars next to the model."""

    model = library_root / "model.safetensors"
    model.write_bytes(b"weights")
    sidecar = library_root / "model.metadata.json"
    sidecar.write_text("{}", encoding="utf-8")
    preview = library_root / "model.preview.webp"
    preview.write_bytes(b"preview")

    await delete_model_artifacts(str(library_root), "model")

    assert not model.exists()
    assert not sidecar.exists()
    assert not preview.exists()


def test_enumerate_model_artifacts_centralized(
    library_root: Path, centralized: Path
):
    model = library_root / "model.safetensors"
    model.write_bytes(b"weights")
    mirror = _mirror_dir(library_root, centralized)
    _write_sidecar(mirror, "model", file_path=model, preview_name="model.preview.png")

    service = PendingDeleteService.__new__(PendingDeleteService)
    artifacts = service._enumerate_model_artifacts(
        str(library_root), "model", ".safetensors"
    )

    assert artifacts == [
        os.path.abspath(str(model)),
        os.path.abspath(str(mirror / "model.metadata.json")),
        os.path.abspath(str(mirror / "model.preview.png")),
    ]


class _RecordingScanner:
    def __init__(self):
        self.calls: List[tuple] = []
        self.model_type = "lora"

    async def update_single_model_cache(self, old_path, new_path, metadata):
        self.calls.append((old_path, new_path, metadata))


class _PassthroughMetadataManager:
    async def save_metadata(self, path: str, metadata):
        await MetadataManager.save_metadata(path, metadata)
        return True


async def _json_metadata_loader(path: str) -> Dict[str, object]:
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


@pytest.mark.asyncio
async def test_rename_model_centralized(library_root: Path, centralized: Path):
    model = library_root / "model.safetensors"
    model.write_bytes(b"weights")
    mirror = _mirror_dir(library_root, centralized)
    _write_sidecar(mirror, "model", file_path=model, preview_name="model.preview.webp")

    service = ModelLifecycleService(
        scanner=_RecordingScanner(),
        metadata_manager=_PassthroughMetadataManager(),
        metadata_loader=_json_metadata_loader,
    )

    result = await service.rename_model(
        file_path=_normalize(model), new_file_name="renamed"
    )

    assert result["success"] is True
    renamed = library_root / "renamed.safetensors"
    assert renamed.exists()
    assert not model.exists()

    # Mirror sidecar/preview renamed with the model; nothing left alongside.
    assert (mirror / "renamed.metadata.json").exists()
    assert (mirror / "renamed.preview.webp").exists()
    assert not (mirror / "model.metadata.json").exists()
    assert not (mirror / "model.preview.webp").exists()
    assert not (library_root / "renamed.metadata.json").exists()
    assert not (library_root / "renamed.preview.webp").exists()

    saved = json.loads((mirror / "renamed.metadata.json").read_text())
    assert saved["file_name"] == "renamed"
    assert saved["file_path"].endswith("renamed.safetensors")
    assert saved["preview_url"].endswith("renamed.preview.webp")
    assert str(mirror).replace(os.sep, "/") in saved["preview_url"]


@pytest.mark.asyncio
async def test_move_model_centralized(
    library_root: Path, centralized: Path, monkeypatch: pytest.MonkeyPatch
):
    source_dir = library_root / "old"
    source_dir.mkdir()
    model = source_dir / "model.safetensors"
    model.write_bytes(b"weights")
    target_dir = library_root / "new"

    old_mirror = _mirror_dir(library_root, centralized, "old")
    _write_sidecar(
        old_mirror, "model", file_path=model, preview_name="model.preview.webp"
    )

    scanner = CheckpointScanner()
    monkeypatch.setattr(
        scanner, "update_single_model_cache", AsyncMock(return_value=True)
    )

    result = await scanner.move_model(_normalize(model), _normalize(target_dir))

    assert result is not None
    moved_model = target_dir / "model.safetensors"
    assert moved_model.exists()
    assert not model.exists()

    new_mirror = _mirror_dir(library_root, centralized, "new")
    assert (new_mirror / "model.metadata.json").exists()
    assert (new_mirror / "model.preview.webp").exists()
    assert not (old_mirror / "model.metadata.json").exists()
    assert not (old_mirror / "model.preview.webp").exists()

    saved = json.loads((new_mirror / "model.metadata.json").read_text())
    assert saved["file_path"].endswith("new/model.safetensors")
    assert saved["preview_url"].endswith("model.preview.webp")
    assert str(new_mirror).replace(os.sep, "/") in saved["preview_url"]


class _FakeCache:
    def __init__(self, raw_data: List[Dict[str, Any]]):
        self.raw_data = raw_data
        self.all_folders: List[str] = []
        self.folders: List[str] = []

    def remove_from_version_index(self, _item) -> None:
        pass

    def rebuild_version_index(self) -> None:
        pass

    async def resort(self) -> None:
        pass


@pytest.mark.asyncio
async def test_rename_known_folder_centralized(
    library_root: Path, centralized: Path, monkeypatch: pytest.MonkeyPatch
):
    # The caller (ModelFileService.rename_folder) has already renamed the
    # model directory on disk when the scanner is asked to re-key records.
    old_dir = library_root / "oldfolder"
    new_dir = library_root / "newfolder"
    new_dir.mkdir()
    model = new_dir / "model.safetensors"
    model.write_bytes(b"weights")

    old_model_path = old_dir / "model.safetensors"
    old_mirror = _mirror_dir(library_root, centralized, "oldfolder")
    _write_sidecar(
        old_mirror, "model", file_path=old_model_path, preview_name="model.preview.webp"
    )

    cache_entry: Dict[str, Any] = {
        "file_path": _normalize(old_model_path),
        "folder": "oldfolder",
        "preview_url": _normalize(old_mirror / "model.preview.webp"),
        "sha256": "",
    }
    scanner = CheckpointScanner()
    scanner._cache = _FakeCache([cache_entry])
    monkeypatch.setattr(scanner, "_persist_current_cache", AsyncMock())

    changed = await scanner.rename_known_folder(
        "oldfolder",
        "newfolder",
        previous_path=str(old_dir),
        new_path=str(new_dir),
    )

    assert changed is True

    new_mirror = _mirror_dir(library_root, centralized, "newfolder")
    assert (new_mirror / "model.metadata.json").exists()
    assert (new_mirror / "model.preview.webp").exists()
    assert not old_mirror.exists()

    saved = json.loads((new_mirror / "model.metadata.json").read_text())
    assert saved["file_path"] == _normalize(model)
    assert saved["preview_url"] == _normalize(new_mirror / "model.preview.webp")

    assert cache_entry["file_path"] == _normalize(model)
    assert cache_entry["folder"] == "newfolder"
    assert cache_entry["preview_url"] == _normalize(
        new_mirror / "model.preview.webp"
    )


@pytest.mark.asyncio
async def test_pending_models_mirror_walk(library_root: Path, centralized: Path):
    model = library_root / "model.safetensors"
    model.write_bytes(b"weights")
    mirror = _mirror_dir(library_root, centralized)
    _write_sidecar(
        mirror,
        "model",
        file_path=model,
        extra={"hash_status": "pending", "sha256": ""},
    )
    # Orphan sidecar: recorded model path is gone and no stem match exists.
    _write_sidecar(
        mirror,
        "ghost",
        file_path=library_root / "ghost.safetensors",
        extra={"hash_status": "pending", "sha256": ""},
    )

    scanner = CheckpointScanner()
    pending = await scanner._find_pending_models_from_filesystem()

    assert len(pending) == 1
    assert pending[0]["file_path"] == _normalize(model)
    assert pending[0]["hash_status"] == "pending"


@pytest.mark.asyncio
async def test_pending_models_mirror_walk_uses_stem_fallback(
    library_root: Path, centralized: Path
):
    """A stale recorded file_path falls back to probing by stem + extension."""

    model = library_root / "model.safetensors"
    model.write_bytes(b"weights")
    mirror = _mirror_dir(library_root, centralized)
    _write_sidecar(
        mirror,
        "model",
        file_path=library_root / "renamed-away.safetensors",
        extra={"hash_status": "pending", "sha256": ""},
    )

    scanner = CheckpointScanner()
    pending = await scanner._find_pending_models_from_filesystem()

    assert len(pending) == 1
    assert pending[0]["file_path"] == _normalize(model)