ComfyUI-Lora-Manager/tests/routes/test_pending_delete_routes.py

"""Tests for the pending-delete undo endpoint.

Covers ``POST /api/lm/undo-delete`` (``py/routes/handlers/
pending_delete_handler.py``) and its route registration
(``py/routes/pending_delete_routes.py``): model/recipe cache restoration,
per-type scanner resolution, restart-safe and rescan-stale undo, tag-count
restoration, broadcast, error responses, malformed input, and exactly-once
registration in both plugin and standalone modes.
"""

from __future__ import annotations

import asyncio
import json
import os
import time
from collections.abc import Iterator
from pathlib import Path
from types import SimpleNamespace
from typing import Any, Dict, List, Optional, Sequence

import pytest
from aiohttp import web

from py.routes.pending_delete_routes import PendingDeleteRoutes
from py.services.pending_delete_service import (
    PENDING_DELETE_DIR_NAME,
    PendingDeleteService,
    _reset_pending_delete_service,
)
from py.utils import settings_paths

UNDO_DELETE_PATH = "/api/lm/undo-delete"


# ---------------------------------------------------------------------------
# Doubles
# ---------------------------------------------------------------------------
class FakeRequest:
    """Request double exposing the async ``json()`` the handler awaits."""

    def __init__(self, *, json_data: Any = None, raise_json_error: bool = False) -> None:
        self._json_data = json_data
        self._raise_json_error = raise_json_error

    async def json(self) -> Any:
        if self._raise_json_error:
            raise ValueError("invalid json body")
        return self._json_data


class FakeHashIndex:
    """Hash index double recording ``add_entry`` calls."""

    def __init__(self) -> None:
        self.entries: List[tuple[Any, ...]] = []

    def add_entry(self, sha256: str, file_path: str, autov3: Optional[str] = None) -> None:
        self.entries.append((sha256, file_path, autov3))


class FakeCache:
    """Cache double exposing the raw_data/version-index surface the undo uses."""

    def __init__(self, items: Optional[Sequence[Dict[str, Any]]] = None) -> None:
        self.raw_data: List[Dict[str, Any]] = list(items or [])
        self.version_index: Dict[Any, Any] = {}
        self.model_id_index: Dict[Any, Any] = {}
        self.rebuild_calls = 0
        self.resort_calls = 0

    def rebuild_version_index(self) -> None:
        self.rebuild_calls += 1
        self.version_index = {}
        self.model_id_index = {}

    async def resort(self) -> None:
        self.resort_calls += 1


class FakeScanner:
    """Scanner double exposing the attributes the undo cache-restore uses."""

    def __init__(
        self,
        root: Path,
        *,
        model_type: str = "lora",
        cache: Optional[FakeCache] = None,
    ) -> None:
        self._root = os.path.abspath(str(root))
        self.model_type = model_type
        self._cache = cache or FakeCache()
        self._hash_index = FakeHashIndex()
        self._tags_count: Dict[str, int] = {}
        self.cache_version = 0
        self.bump_calls = 0
        self.persist_calls = 0

    def get_model_roots(self) -> List[str]:
        return [self._root]

    def _find_root_for_file(self, file_path: Optional[str]) -> Optional[str]:
        if not file_path:
            return None
        normalized = os.path.abspath(os.path.normpath(file_path))
        if normalized == self._root or normalized.startswith(self._root + os.sep):
            return self._root
        return None

    async def get_cached_data(self, force_refresh: bool = False) -> FakeCache:
        return self._cache

    def bump_cache_version(self) -> None:
        self.cache_version += 1
        self.bump_calls += 1

    async def _persist_current_cache(self) -> None:
        self.persist_calls += 1


class FakeRecipeScanner:
    """Recipe scanner double mirroring add_recipe's tolerant path-map read."""

    def __init__(self) -> None:
        self.added: List[Dict[str, Any]] = []
        self._json_path_map: Dict[str, str] = {}

    async def add_recipe(self, recipe_data: Dict[str, Any]) -> None:
        # The real scanner only READS _json_path_map here (the row may carry an
        # empty json_path until the forced frontend refresh) - never crashes.
        recipe_id = str(recipe_data.get("id", ""))
        self._json_path_map.get(recipe_id, "")
        self.added.append(dict(recipe_data))

    def force_refresh(self) -> None:
        """Simulate ``window.recipeManager.loadRecipes(true)`` rebuilding the map."""
        for recipe in self.added:
            recipe_id = str(recipe.get("id", ""))
            file_path = recipe.get("file_path")
            if recipe_id and file_path:
                self._json_path_map[recipe_id] = str(file_path)


class _DummyRoutes:
    @staticmethod
    def setup_routes(app_: web.Application, **kwargs: Any) -> None:
        return None


class _DummyWSManager:
    async def handle_connection(self, request):  # pragma: no cover - interface stub
        return None

    async def handle_download_connection(self, request):  # pragma: no cover - interface stub
        return None

    async def handle_init_connection(self, request):  # pragma: no cover - interface stub
        return None


# ---------------------------------------------------------------------------
# Fixtures / helpers
# ---------------------------------------------------------------------------
@pytest.fixture(autouse=True)
def _reset_service_singleton() -> Iterator[None]:
    """Reset the pending-delete singleton before and after each test."""
    _reset_pending_delete_service()
    yield
    _reset_pending_delete_service()


@pytest.fixture(autouse=True)
def _stub_scanner_registry(monkeypatch: pytest.MonkeyPatch) -> None:
    """Prevent scanner resolution from instantiating real singletons."""
    from py.services.service_registry import ServiceRegistry

    async def _none(*_args: Any, **_kwargs: Any) -> None:
        return None

    monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _none)
    monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _none)
    monkeypatch.setattr(ServiceRegistry, "get_embedding_scanner", _none)
    monkeypatch.setattr(ServiceRegistry, "get_recipe_scanner", _none)


async def _register_scanners(
    monkeypatch: pytest.MonkeyPatch,
    *,
    lora: Any = None,
    checkpoint: Any = None,
    embedding: Any = None,
    recipe: Any = None,
) -> None:
    """Point the ServiceRegistry getters at per-test scanner doubles."""
    from py.services.service_registry import ServiceRegistry

    def _make(scanner: Any):
        async def _getter(*_args: Any, **_kwargs: Any) -> Any:
            return scanner

        return _getter

    monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _make(lora))
    monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _make(checkpoint))
    monkeypatch.setattr(ServiceRegistry, "get_embedding_scanner", _make(embedding))
    monkeypatch.setattr(ServiceRegistry, "get_recipe_scanner", _make(recipe))


def _make_undo_request(**json_data: Any) -> FakeRequest:
    return FakeRequest(json_data=json_data)


def _json_payload(response: web.Response) -> Dict[str, Any]:
    """Decode the JSON body of a web.Response, asserting it is not null."""
    text = response.text
    assert text is not None
    return json.loads(text)


def _write_batch_manifest(
    batch_dir: Path,
    *,
    batch_id: str,
    kind: str,
    expires_at: int,
    entries: Sequence[Dict[str, Any]],
    model_type: Optional[str] = None,
    state: str = "staged",
    model_snapshot: Any = None,
    recipe_snapshot: Any = None,
) -> None:
    """Write a manifest.json with the shape the service reads."""
    manifest: Dict[str, Any] = {
        "batch_id": batch_id,
        "kind": kind,
        "model_type": model_type if kind == "model" else None,
        "state": state,
        "expires_at": int(expires_at),
        "entries": list(entries),
        "model_snapshot": model_snapshot if kind == "model" else None,
        "recipe_snapshot": recipe_snapshot if kind == "recipe" else None,
    }
    (batch_dir / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")


def _undo_delete_routes(app: web.Application) -> List[Any]:
    """The POST /api/lm/undo-delete routes currently registered on *app*."""
    return [
        route
        for route in app.router.routes()
        if route.method == "POST"
        and getattr(route.resource, "canonical", "") == UNDO_DELETE_PATH
    ]


async def _drain_broadcast() -> None:
    """Yield to the loop so the fire-and-forget broadcast task can run."""
    for _ in range(10):
        await asyncio.sleep(0)


# (a) undo of a staged model batch (loras) -> files + cache + hash + broadcast
# ---------------------------------------------------------------------------
async def test_a_undo_staged_model_batch_restores_files_cache_and_broadcasts(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "loras"
    root.mkdir()
    model = root / "model.safetensors"
    model.write_bytes(b"model-data")
    snapshot = {
        "file_path": str(model),
        "sha256": "a" * 64,
        "tags": ["alpha"],
        "model_name": "model",
    }

    scanner = FakeScanner(root, model_type="lora")
    scanner._cache.raw_data = [dict(snapshot)]
    scanner._tags_count = {"alpha": 1}
    await _register_scanners(monkeypatch, lora=scanner)

    service = await PendingDeleteService.get_instance()
    batch_id = await service.stage_model_delete(
        scanner=scanner,
        target_dir=str(root),
        file_name="model",
        main_extension=".safetensors",
        original_file_path=str(model),
        cached_entry=dict(snapshot),
    )
    assert batch_id is not None

    # Simulate the delete-time cache mutation: entry removed, tags decremented.
    scanner._cache.raw_data = []
    scanner._tags_count = {}

    sent: List[Dict[str, Any]] = []

    async def fake_broadcast(data: Dict[str, Any]) -> None:
        sent.append(data)

    monkeypatch.setattr("py.services.websocket_manager.ws_manager.broadcast", fake_broadcast)

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=batch_id)  # pyright: ignore[reportArgumentType]
    )

    assert response.status == 200
    payload = _json_payload(response)
    assert payload["success"] is True
    assert payload["kind"] == "model"
    assert payload["restored"] == [str(model)]

    # Files restored to the original path.
    assert model.read_bytes() == b"model-data"
    # Cache entry restored exactly once.
    assert len(scanner._cache.raw_data) == 1
    assert scanner._cache.raw_data[0]["file_path"] == str(model)
    # Version-index rebuild + resort + persist + cache bump all ran.
    assert scanner._cache.rebuild_calls >= 1
    assert scanner._cache.resort_calls >= 1
    assert scanner.persist_calls >= 1
    assert scanner.bump_calls >= 1
    # Hash index has the path.
    assert any(entry[1] == str(model) for entry in scanner._hash_index.entries)
    # Tag counts re-incremented.
    assert scanner._tags_count == {"alpha": 1}
    # Broadcast fired exactly once with models_changed.
    await _drain_broadcast()
    assert sent == [{"type": "models_changed"}]


# ---------------------------------------------------------------------------
# (b) undo of a staged CHECKPOINT batch -> resolves the checkpoint scanner
# ---------------------------------------------------------------------------
async def test_b_undo_checkpoint_batch_updates_checkpoint_cache_only(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    ckpt_root = tmp_path / "checkpoints"
    ckpt_root.mkdir()
    model = ckpt_root / "model.safetensors"
    model.write_bytes(b"ckpt-data")
    snapshot = {
        "file_path": str(model),
        "sha256": "b" * 64,
        "tags": [],
        "model_name": "ckpt",
    }

    lora_scanner = FakeScanner(tmp_path / "loras", model_type="lora")
    ckpt_scanner = FakeScanner(ckpt_root, model_type="checkpoint")
    lora_scanner._cache.raw_data = [{"file_path": "/loras/other.safetensors", "sha256": "x"}]
    ckpt_scanner._cache.raw_data = [dict(snapshot)]
    await _register_scanners(monkeypatch, lora=lora_scanner, checkpoint=ckpt_scanner)

    service = await PendingDeleteService.get_instance()
    batch_id = await service.stage_model_delete(
        scanner=ckpt_scanner,
        target_dir=str(ckpt_root),
        file_name="model",
        main_extension=".safetensors",
        original_file_path=str(model),
        cached_entry=dict(snapshot),
    )
    assert batch_id is not None
    # Simulate the delete-time cache mutation on the checkpoint cache.
    ckpt_scanner._cache.raw_data = []

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=batch_id)  # pyright: ignore[reportArgumentType]
    )
    assert response.status == 200

    # The CHECKPOINT cache was updated, not the lora cache.
    assert len(ckpt_scanner._cache.raw_data) == 1
    assert ckpt_scanner._cache.raw_data[0]["file_path"] == str(model)
    assert len(lora_scanner._cache.raw_data) == 1  # untouched


# ---------------------------------------------------------------------------
# (c) undo of a staged recipe batch -> files restored + add_recipe called
# ---------------------------------------------------------------------------
async def test_c_undo_recipe_batch_restores_files_and_calls_add_recipe(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    recipe_json = tmp_path / "my_recipe.recipe.json"
    recipe_data = {"id": "r1", "name": "Recipe One", "file_path": str(recipe_json)}
    recipe_json.write_text(json.dumps(recipe_data), encoding="utf-8")

    recipe_scanner = FakeRecipeScanner()
    await _register_scanners(monkeypatch, recipe=recipe_scanner)

    service = await PendingDeleteService.get_instance()
    batch_id = await service.stage_recipe_delete(
        recipe_json_path=str(recipe_json),
        image_path=None,
        recipe_data=dict(recipe_data),
    )
    assert batch_id is not None
    # The todo-4 caller removes the originals after staging.
    recipe_json.unlink()

    sent: List[Dict[str, Any]] = []

    async def fake_broadcast(data: Dict[str, Any]) -> None:
        sent.append(data)

    monkeypatch.setattr("py.services.websocket_manager.ws_manager.broadcast", fake_broadcast)

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=batch_id)  # pyright: ignore[reportArgumentType]
    )

    assert response.status == 200
    payload = _json_payload(response)
    assert payload["success"] is True
    assert payload["kind"] == "recipe"
    assert payload["restored"] == [str(recipe_json)]

    # Files restored and the recipe re-added to the scanner.
    assert recipe_json.read_text(encoding="utf-8") == json.dumps(recipe_data)
    assert recipe_scanner.added == [recipe_data]
    # Recipe undo is client-refresh only - no models_changed broadcast.
    await _drain_broadcast()
    assert sent == []


# ---------------------------------------------------------------------------
# (d) unknown / expired / occupied batch -> 404 responses
# ---------------------------------------------------------------------------
async def test_d1_undo_unknown_batch_returns_404(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    await _register_scanners(
        monkeypatch, lora=FakeScanner(tmp_path / "loras", model_type="lora")
    )

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id="f" * 32)  # pyright: ignore[reportArgumentType]
    )

    assert response.status == 404
    payload = _json_payload(response)
    assert payload["success"] is False
    assert "batch" in payload["error"].lower()


async def test_d2_undo_expired_batch_returns_404(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "loras"
    root.mkdir()
    model = root / "model.safetensors"
    model.write_bytes(b"data")

    scanner = FakeScanner(root, model_type="lora")
    await _register_scanners(monkeypatch, lora=scanner)

    service = await PendingDeleteService.get_instance()

    # Isolate the expiry check from the opportunistic purge.
    async def _no_purge() -> int:
        return 0

    monkeypatch.setattr(service, "purge_expired", _no_purge)

    batch_id = await service.stage_model_delete(
        scanner=scanner,
        target_dir=str(root),
        file_name="model",
        main_extension=".safetensors",
        original_file_path=str(model),
        cached_entry={"file_path": str(model)},
    )
    assert batch_id is not None
    batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id
    manifest_path = batch_dir / "manifest.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    manifest["expires_at"] = int(time.time()) - 10
    manifest_path.write_text(json.dumps(manifest), encoding="utf-8")

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=batch_id)  # pyright: ignore[reportArgumentType]
    )
    assert response.status == 404
    assert "expired" in _json_payload(response)["error"].lower()


async def test_d3_undo_occupied_path_returns_404(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "loras"
    root.mkdir()
    model = root / "model.safetensors"
    model.write_bytes(b"data")

    scanner = FakeScanner(root, model_type="lora")
    await _register_scanners(monkeypatch, lora=scanner)

    service = await PendingDeleteService.get_instance()
    batch_id = await service.stage_model_delete(
        scanner=scanner,
        target_dir=str(root),
        file_name="model",
        main_extension=".safetensors",
        original_file_path=str(model),
        cached_entry={"file_path": str(model)},
    )
    assert batch_id is not None
    # A re-download now occupies the original path.
    model.write_bytes(b"new-file")

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=batch_id)  # pyright: ignore[reportArgumentType]
    )
    assert response.status == 404
    assert "occupied" in _json_payload(response)["error"].lower()


# ---------------------------------------------------------------------------
# (d-adjacent) malformed input -> 400
# ---------------------------------------------------------------------------
async def test_d4_undo_malformed_body_returns_400(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    handler = PendingDeleteHandler()

    # Missing batch_id.
    response = await handler.undo_delete(_make_undo_request())  # pyright: ignore[reportArgumentType]
    assert response.status == 400
    assert "batch_id" in _json_payload(response)["error"].lower()

    # Non-dict JSON body.
    response = await handler.undo_delete(FakeRequest(json_data="not-a-dict"))  # pyright: ignore[reportArgumentType]
    assert response.status == 400

    # Invalid JSON body.
    response = await handler.undo_delete(FakeRequest(raise_json_error=True))  # pyright: ignore[reportArgumentType]
    assert response.status == 400

    # (VAL-1) path-traversal style batch_id -> 400.
    response = await handler.undo_delete(_make_undo_request(batch_id="../evil"))  # pyright: ignore[reportArgumentType]
    assert response.status == 400
    assert "Invalid batch_id" in _json_payload(response)["error"]

    # (VAL-2) non-hex batch_id -> 400.
    response = await handler.undo_delete(_make_undo_request(batch_id="not-a-hex-id"))  # pyright: ignore[reportArgumentType]
    assert response.status == 400
    assert "Invalid batch_id" in _json_payload(response)["error"]


# ---------------------------------------------------------------------------
# (e) route registered exactly ONCE in both modes
# ---------------------------------------------------------------------------
def test_e1_routes_class_registers_undo_route_exactly_once() -> None:
    app = web.Application()
    PendingDeleteRoutes.setup_routes(app)

    assert len(_undo_delete_routes(app)) == 1


def test_e2_plugin_mode_registers_undo_route_exactly_once(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    from py import lora_manager

    app = web.Application()
    app._handler_args = {"max_field_size": 1024}
    monkeypatch.setattr(lora_manager.PromptServer, "instance", SimpleNamespace(app=app))

    monkeypatch.setattr(lora_manager, "register_default_model_types", lambda: None)
    monkeypatch.setattr(
        lora_manager.ModelServiceFactory, "setup_all_routes", lambda app_: None
    )
    monkeypatch.setattr(
        lora_manager.ModelServiceFactory, "get_registered_types", lambda: ["dummy"]
    )
    for name in (
        "StatsRoutes",
        "RecipeRoutes",
        "UpdateRoutes",
        "MiscRoutes",
        "ExampleImagesRoutes",
        "PreviewRoutes",
    ):
        monkeypatch.setattr(lora_manager, name, _DummyRoutes)
    monkeypatch.setattr(lora_manager, "ws_manager", _DummyWSManager())
    monkeypatch.setattr(
        lora_manager.settings,
        "get",
        lambda key, default=None: str(tmp_path) if key == "example_images_path" else default,
    )
    monkeypatch.setattr(app.router, "add_static", lambda *a, **k: SimpleNamespace())

    lora_manager.LoraManager.add_routes()

    assert len(_undo_delete_routes(app)) == 1


def test_e3_standalone_mode_registers_undo_route_exactly_once(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    import standalone as standalone_module

    app = web.Application()
    locales_dir = tmp_path / "locales"
    locales_dir.mkdir()
    static_dir = tmp_path / "static"
    static_dir.mkdir()
    monkeypatch.setattr(
        standalone_module,
        "config",
        SimpleNamespace(i18n_path=str(locales_dir), static_path=str(static_dir)),
    )

    import py.services.model_service_factory as factory_module

    monkeypatch.setattr(factory_module, "register_default_model_types", lambda: None)
    monkeypatch.setattr(
        factory_module.ModelServiceFactory,
        "setup_all_routes",
        classmethod(lambda cls, app_arg: None),
    )
    monkeypatch.setattr("py.routes.recipe_routes.RecipeRoutes", _DummyRoutes)
    monkeypatch.setattr("py.routes.update_routes.UpdateRoutes", _DummyRoutes)
    monkeypatch.setattr("py.routes.misc_routes.MiscRoutes", _DummyRoutes)
    monkeypatch.setattr("py.routes.stats_routes.StatsRoutes", _DummyRoutes)
    monkeypatch.setattr("py.routes.example_images_routes.ExampleImagesRoutes", _DummyRoutes)
    monkeypatch.setattr("py.routes.preview_routes.PreviewRoutes", _DummyRoutes)

    async def _noop_ws_handler(request: Any) -> web.Response:
        return web.Response(status=204)

    ws_stub = SimpleNamespace(
        handle_connection=_noop_ws_handler,
        handle_download_connection=_noop_ws_handler,
        handle_init_connection=_noop_ws_handler,
    )
    monkeypatch.setattr("py.services.websocket_manager.ws_manager", ws_stub)

    server = SimpleNamespace(app=app)
    standalone_module.StandaloneLoraManager.add_routes(server)

    assert len(_undo_delete_routes(app)) == 1


# ---------------------------------------------------------------------------
# (f) RESTART-SAFE UNDO (T8): staged batch on disk, no in-process timer
# ---------------------------------------------------------------------------
async def test_f_restart_safe_undo_restores_files_and_cache(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "loras"
    root.mkdir()
    model = root / "model.safetensors"
    # NOTE: the original file does NOT exist - it lives in the staging batch
    # (a delete moved it away before the restart).
    snapshot = {"file_path": str(model), "sha256": "c" * 64, "tags": ["beta"]}

    # Simulate a post-restart state: the batch + manifest exist on disk but
    # nothing was staged in-process (no timer task, no _known_roots entry).
    restart_batch_id = "b" * 32
    batch_dir = root / PENDING_DELETE_DIR_NAME / restart_batch_id
    batch_dir.mkdir(parents=True)
    (batch_dir / "model.safetensors").write_bytes(b"survived-restart")
    _write_batch_manifest(
        batch_dir,
        batch_id=restart_batch_id,
        kind="model",
        model_type="loras",
        expires_at=int(time.time()) + 3600,
        entries=[
            {
                "staged": str(batch_dir / "model.safetensors"),
                "original": str(model),
                "restored": False,
            }
        ],
        model_snapshot=dict(snapshot),
    )

    scanner = FakeScanner(root, model_type="lora")
    await _register_scanners(monkeypatch, lora=scanner)

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=restart_batch_id)  # pyright: ignore[reportArgumentType]
    )

    assert response.status == 200
    assert model.read_bytes() == b"survived-restart"
    assert not batch_dir.exists()
    assert len(scanner._cache.raw_data) == 1
    assert scanner._cache.raw_data[0]["file_path"] == str(model)
    assert scanner._tags_count == {"beta": 1}


# ---------------------------------------------------------------------------
# (g) RESCAN STALENESS (T7): stale raw_data entry replaced by the snapshot
# ---------------------------------------------------------------------------
async def test_g_rescan_stale_cache_entry_replaced_by_snapshot(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "loras"
    root.mkdir()
    model = root / "model.safetensors"
    model.write_bytes(b"model-data")
    snapshot = {
        "file_path": str(model),
        "sha256": "d" * 64,
        "tags": ["gamma"],
        "model_name": "model",
    }

    scanner = FakeScanner(root, model_type="lora")
    await _register_scanners(monkeypatch, lora=scanner)

    service = await PendingDeleteService.get_instance()
    batch_id = await service.stage_model_delete(
        scanner=scanner,
        target_dir=str(root),
        file_name="model",
        main_extension=".safetensors",
        original_file_path=str(model),
        cached_entry=dict(snapshot),
    )
    assert batch_id is not None

    # A rescan between delete and undo re-added a STALE entry for the same path.
    stale = {
        "file_path": str(model),
        "sha256": "z" * 64,
        "tags": ["stale-tag"],
        "model_name": "stale",
    }
    scanner._cache.raw_data = [stale]

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=batch_id)  # pyright: ignore[reportArgumentType]
    )
    assert response.status == 200

    # Exactly one entry remains and it equals the snapshot.
    assert len(scanner._cache.raw_data) == 1
    assert scanner._cache.raw_data[0] == snapshot
    assert scanner._tags_count == {"gamma": 1}


# ---------------------------------------------------------------------------
# (h) RECIPE RE-DELETE AFTER UNDO (T9): path map rebuilt by forced refresh
# ---------------------------------------------------------------------------
async def test_h_recipe_redo_delete_after_undo_and_refresh_succeeds(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    recipe_json = tmp_path / "my_recipe.recipe.json"
    recipe_data = {"id": "r1", "name": "Recipe One", "file_path": str(recipe_json)}
    recipe_json.write_text(json.dumps(recipe_data), encoding="utf-8")

    recipe_scanner = FakeRecipeScanner()
    await _register_scanners(monkeypatch, recipe=recipe_scanner)

    service = await PendingDeleteService.get_instance()
    batch_id = await service.stage_recipe_delete(
        recipe_json_path=str(recipe_json),
        image_path=None,
        recipe_data=dict(recipe_data),
    )
    assert batch_id is not None
    recipe_json.unlink()

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    handler = PendingDeleteHandler()
    response = await handler.undo_delete(_make_undo_request(batch_id=batch_id))  # pyright: ignore[reportArgumentType]
    assert response.status == 200
    assert recipe_json.exists()
    assert recipe_scanner.added == [recipe_data]

    # Forced frontend refresh rebuilds the scanner path map.
    recipe_scanner.force_refresh()
    assert recipe_scanner._json_path_map.get("r1") == str(recipe_json)

    # Immediately deleting the same recipe again succeeds.
    batch_id2 = await service.stage_recipe_delete(
        recipe_json_path=str(recipe_json),
        image_path=None,
        recipe_data=dict(recipe_data),
    )
    assert batch_id2 is not None
    recipe_json.unlink()
    response2 = await handler.undo_delete(_make_undo_request(batch_id=batch_id2))  # pyright: ignore[reportArgumentType]
    assert response2.status == 200
    assert recipe_json.exists()


# ---------------------------------------------------------------------------
# (i) recipe undo WITHOUT prior refresh still restores files (no crash)
# ---------------------------------------------------------------------------
async def test_i_recipe_undo_without_refresh_still_restores_files(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    recipe_json = tmp_path / "no_refresh.recipe.json"
    recipe_data = {"id": "r2", "name": "No Refresh"}
    recipe_json.write_text(json.dumps(recipe_data), encoding="utf-8")

    recipe_scanner = FakeRecipeScanner()
    await _register_scanners(monkeypatch, recipe=recipe_scanner)

    service = await PendingDeleteService.get_instance()
    batch_id = await service.stage_recipe_delete(
        recipe_json_path=str(recipe_json),
        image_path=None,
        recipe_data=dict(recipe_data),
    )
    assert batch_id is not None
    recipe_json.unlink()

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=batch_id)  # pyright: ignore[reportArgumentType]
    )

    assert response.status == 200
    assert recipe_json.read_text(encoding="utf-8") == json.dumps(recipe_data)
    assert recipe_scanner.added == [recipe_data]


# ---------------------------------------------------------------------------
# (j) TAG COUNTS: bulk-delete 2 tagged models, undo restores pre-delete counts
# ---------------------------------------------------------------------------
async def test_j_undo_restores_tag_counts_after_bulk_delete(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "loras"
    root.mkdir()
    model_a = root / "alpha.safetensors"
    model_a.write_bytes(b"alpha")
    model_b = root / "beta.safetensors"
    model_b.write_bytes(b"beta")

    snapshot_a = {"file_path": str(model_a), "sha256": "e" * 64, "tags": ["alpha"]}
    snapshot_b = {"file_path": str(model_b), "sha256": "f" * 64, "tags": ["beta"]}

    scanner = FakeScanner(root, model_type="lora")
    await _register_scanners(monkeypatch, lora=scanner)

    service = await PendingDeleteService.get_instance()
    batch_a = await service.stage_model_delete(
        scanner=scanner,
        target_dir=str(root),
        file_name="alpha",
        main_extension=".safetensors",
        original_file_path=str(model_a),
        cached_entry=dict(snapshot_a),
    )
    batch_b = await service.stage_model_delete(
        scanner=scanner,
        target_dir=str(root),
        file_name="beta",
        main_extension=".safetensors",
        original_file_path=str(model_b),
        cached_entry=dict(snapshot_b),
    )
    assert batch_a is not None
    assert batch_b is not None

    # Simulate the bulk-delete cache mutation (mirror of
    # _batch_update_cache_for_deleted_models): entries removed, tags decremented.
    scanner._cache.raw_data = [dict(snapshot_a), dict(snapshot_b)]
    scanner._tags_count = {"alpha": 1, "beta": 1}
    for model in (dict(snapshot_a), dict(snapshot_b)):
        for tag in model.get("tags", []):
            if tag in scanner._tags_count:
                scanner._tags_count[tag] = max(0, scanner._tags_count[tag] - 1)
                if scanner._tags_count[tag] == 0:
                    del scanner._tags_count[tag]
    scanner._cache.raw_data = []
    pre_delete_counts = {"alpha": 1, "beta": 1}

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    handler = PendingDeleteHandler()
    # Sequential undo of the (unmerged) constituent batches - the batch_ids
    # fallback path used by the frontend for cross-volume bulks.
    assert (await handler.undo_delete(_make_undo_request(batch_id=batch_a))).status == 200  # pyright: ignore[reportArgumentType]
    assert (await handler.undo_delete(_make_undo_request(batch_id=batch_b))).status == 200  # pyright: ignore[reportArgumentType]

    assert scanner._tags_count == pre_delete_counts
    assert {item["file_path"] for item in scanner._cache.raw_data} == {
        str(model_a),
        str(model_b),
    }
    assert model_a.read_bytes() == b"alpha"
    assert model_b.read_bytes() == b"beta"


# ---------------------------------------------------------------------------
# (k) EMBEDDINGS UNDO: resolves the embedding scanner, not the lora scanner
# ---------------------------------------------------------------------------
async def test_k_undo_embeddings_batch_updates_embedding_cache_only(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    emb_root = tmp_path / "embeddings"
    emb_root.mkdir()
    model = emb_root / "my_embedding.safetensors"
    model.write_bytes(b"emb-data")
    snapshot = {
        "file_path": str(model),
        "sha256": "g" * 64,
        "tags": [],
        "model_name": "my_embedding",
    }

    lora_scanner = FakeScanner(tmp_path / "loras", model_type="lora")
    emb_scanner = FakeScanner(emb_root, model_type="embedding")
    lora_scanner._cache.raw_data = [{"file_path": "/loras/other.safetensors", "sha256": "x"}]
    emb_scanner._cache.raw_data = [dict(snapshot)]
    await _register_scanners(monkeypatch, lora=lora_scanner, embedding=emb_scanner)

    service = await PendingDeleteService.get_instance()
    batch_id = await service.stage_model_delete(
        scanner=emb_scanner,
        target_dir=str(emb_root),
        file_name="my_embedding",
        main_extension=".safetensors",
        original_file_path=str(model),
        cached_entry=dict(snapshot),
    )
    assert batch_id is not None
    emb_scanner._cache.raw_data = []

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=batch_id)  # pyright: ignore[reportArgumentType]
    )
    assert response.status == 200

    # The EMBEDDINGS cache was updated, not the lora cache.
    assert len(emb_scanner._cache.raw_data) == 1
    assert emb_scanner._cache.raw_data[0]["file_path"] == str(model)
    assert len(lora_scanner._cache.raw_data) == 1  # untouched
    assert model.read_bytes() == b"emb-data"


# ---------------------------------------------------------------------------
# (BULK-UNDO) MERGED bulk undo restores ALL cache entries (not just the winner)
# ---------------------------------------------------------------------------
async def test_bulk_undo_after_merge_restores_all_cache_entries(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "loras"
    root.mkdir()
    model_a = root / "alpha.safetensors"
    model_a.write_bytes(b"alpha")
    model_b = root / "beta.safetensors"
    model_b.write_bytes(b"beta")

    snapshot_a = {
        "file_path": str(model_a),
        "sha256": "e" * 64,
        "tags": ["alpha"],
        "model_name": "alpha",
    }
    snapshot_b = {
        "file_path": str(model_b),
        "sha256": "f" * 64,
        "tags": ["beta"],
        "model_name": "beta",
    }

    scanner = FakeScanner(root, model_type="lora")
    await _register_scanners(monkeypatch, lora=scanner)

    service = await PendingDeleteService.get_instance()
    batch_a = await service.stage_model_delete(
        scanner=scanner,
        target_dir=str(root),
        file_name="alpha",
        main_extension=".safetensors",
        original_file_path=str(model_a),
        cached_entry=dict(snapshot_a),
    )
    batch_b = await service.stage_model_delete(
        scanner=scanner,
        target_dir=str(root),
        file_name="beta",
        main_extension=".safetensors",
        original_file_path=str(model_b),
        cached_entry=dict(snapshot_b),
    )
    assert batch_a is not None
    assert batch_b is not None
    merged = await service.merge_batches([batch_a, batch_b])
    assert merged == batch_a

    # Simulate the bulk-delete cache mutation (entries removed, tags cleared).
    scanner._cache.raw_data = []
    scanner._tags_count = {}
    scanner._hash_index = FakeHashIndex()

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=merged)  # pyright: ignore[reportArgumentType]
    )
    assert response.status == 200

    # BOTH entries restored - the loser's file_path present in raw_data.
    assert {item["file_path"] for item in scanner._cache.raw_data} == {
        str(model_a),
        str(model_b),
    }
    # Hash index carries both paths.
    assert {entry[1] for entry in scanner._hash_index.entries} == {
        str(model_a),
        str(model_b),
    }
    # Tag counts match pre-delete counts for tagged models.
    assert scanner._tags_count == {"alpha": 1, "beta": 1}
    # Files restored byte-identical.
    assert model_a.read_bytes() == b"alpha"
    assert model_b.read_bytes() == b"beta"


# ---------------------------------------------------------------------------
# (BC-1) backward compat: entries WITHOUT snapshots + top-level model_snapshot
# ---------------------------------------------------------------------------
async def test_bc1_old_format_manifest_falls_back_to_top_level_snapshot(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "loras"
    root.mkdir()
    model = root / "model.safetensors"
    snapshot = {
        "file_path": str(model),
        "sha256": "c" * 64,
        "tags": ["beta"],
        "model_name": "model",
    }

    # Old-format manifest: entries carry NO snapshot; only the top-level
    # model_snapshot exists (pre-F3 single-delete manifests).
    batch_id = "a" * 32
    batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id
    batch_dir.mkdir(parents=True)
    (batch_dir / "model.safetensors").write_bytes(b"old-format-data")
    _write_batch_manifest(
        batch_dir,
        batch_id=batch_id,
        kind="model",
        model_type="loras",
        expires_at=int(time.time()) + 3600,
        entries=[
            {
                "staged": str(batch_dir / "model.safetensors"),
                "original": str(model),
                "restored": False,
            }
        ],
        model_snapshot=dict(snapshot),
    )

    scanner = FakeScanner(root, model_type="lora")
    await _register_scanners(monkeypatch, lora=scanner)

    from py.routes.handlers.pending_delete_handler import PendingDeleteHandler

    response = await PendingDeleteHandler().undo_delete(
        _make_undo_request(batch_id=batch_id)  # pyright: ignore[reportArgumentType]
    )
    assert response.status == 200

    # The top-level snapshot entry is still restored exactly.
    assert len(scanner._cache.raw_data) == 1
    assert scanner._cache.raw_data[0] == snapshot
    assert scanner._tags_count == {"beta": 1}
    assert model.read_bytes() == b"old-format-data"