invoke-ai--invokeai/tests/app/services/model_install/test_missing_models.py

"""
Tests for missing model detection (_scan_for_missing_models) and bulk deletion.
"""

import gc
from pathlib import Path

import pytest

from invokeai.app.services.config import InvokeAIAppConfig
from invokeai.app.services.model_install import ModelInstallServiceBase
from invokeai.app.services.model_records import UnknownModelException
from invokeai.backend.model_manager.configs.textual_inversion import TI_File_SD1_Config
from invokeai.backend.model_manager.taxonomy import (
    BaseModelType,
    ModelFormat,
    ModelSourceType,
    ModelType,
)
from tests.backend.model_manager.model_manager_fixtures import *  # noqa F403


class TestScanForMissingModels:
    """Tests for ModelInstallService._scan_for_missing_models()."""

    def test_no_missing_models(
        self, mm2_installer: ModelInstallServiceBase, embedding_file: Path, mm2_app_config: InvokeAIAppConfig
    ) -> None:
        """When all registered models exist on disk, _scan_for_missing_models returns an empty list."""
        mm2_installer.register_path(embedding_file)
        missing = mm2_installer._scan_for_missing_models()
        assert len(missing) == 0

    def test_detects_missing_model(
        self, mm2_installer: ModelInstallServiceBase, embedding_file: Path, mm2_app_config: InvokeAIAppConfig
    ) -> None:
        """A model whose path does not exist on disk is reported as missing."""
        # Register a real model first, then add a fake one with a non-existent path
        mm2_installer.register_path(embedding_file)

        fake_config = TI_File_SD1_Config(
            key="missing-model-key-1",
            path="/nonexistent/path/missing_model.safetensors",
            name="MissingModel",
            base=BaseModelType.StableDiffusion1,
            type=ModelType.TextualInversion,
            format=ModelFormat.EmbeddingFile,
            hash="FAKEHASH1",
            file_size=1024,
            source="test/source",
            source_type=ModelSourceType.Path,
        )
        mm2_installer.record_store.add_model(fake_config)

        missing = mm2_installer._scan_for_missing_models()
        assert len(missing) == 1
        assert missing[0].key == "missing-model-key-1"

    def test_mix_of_existing_and_missing(
        self,
        mm2_installer: ModelInstallServiceBase,
        embedding_file: Path,
        diffusers_dir: Path,
        mm2_app_config: InvokeAIAppConfig,
    ) -> None:
        """With multiple models, only the ones with missing files are returned."""
        key_existing = mm2_installer.register_path(embedding_file)
        mm2_installer.register_path(diffusers_dir)

        # Add two models with non-existent paths
        fake1 = TI_File_SD1_Config(
            key="missing-key-1",
            path="/nonexistent/missing1.safetensors",
            name="Missing1",
            base=BaseModelType.StableDiffusion1,
            type=ModelType.TextualInversion,
            format=ModelFormat.EmbeddingFile,
            hash="FAKEHASH_A",
            file_size=1024,
            source="test/source1",
            source_type=ModelSourceType.Path,
        )
        fake2 = TI_File_SD1_Config(
            key="missing-key-2",
            path="/nonexistent/missing2.safetensors",
            name="Missing2",
            base=BaseModelType.StableDiffusion1,
            type=ModelType.TextualInversion,
            format=ModelFormat.EmbeddingFile,
            hash="FAKEHASH_B",
            file_size=2048,
            source="test/source2",
            source_type=ModelSourceType.Path,
        )
        mm2_installer.record_store.add_model(fake1)
        mm2_installer.record_store.add_model(fake2)

        missing = mm2_installer._scan_for_missing_models()
        missing_keys = {m.key for m in missing}
        assert len(missing) == 2
        assert "missing-key-1" in missing_keys
        assert "missing-key-2" in missing_keys
        assert key_existing not in missing_keys

    def test_empty_store_returns_empty(self, mm2_installer: ModelInstallServiceBase) -> None:
        """With no models registered, _scan_for_missing_models returns an empty list."""
        missing = mm2_installer._scan_for_missing_models()
        assert len(missing) == 0


class TestBulkDelete:
    """Tests for bulk model deletion."""

    def test_delete_installed_model(
        self, mm2_installer: ModelInstallServiceBase, embedding_file: Path, mm2_app_config: InvokeAIAppConfig
    ) -> None:
        """Deleting an installed model removes it from the store and disk."""
        key = mm2_installer.install_path(embedding_file)
        record = mm2_installer.record_store.get_model(key)
        model_path = mm2_app_config.models_path / record.path
        assert model_path.exists()
        assert mm2_installer.record_store.exists(key)

        gc.collect()
        mm2_installer.delete(key)

        with pytest.raises(UnknownModelException):
            mm2_installer.record_store.get_model(key)

    def test_unregister_missing_model(
        self, mm2_installer: ModelInstallServiceBase, mm2_app_config: InvokeAIAppConfig
    ) -> None:
        """Unregistering a model whose file is missing removes it from the DB."""
        fake_config = TI_File_SD1_Config(
            key="missing-to-delete",
            path="/nonexistent/path/gone.safetensors",
            name="GoneModel",
            base=BaseModelType.StableDiffusion1,
            type=ModelType.TextualInversion,
            format=ModelFormat.EmbeddingFile,
            hash="FAKEHASH_GONE",
            file_size=1024,
            source="test/source",
            source_type=ModelSourceType.Path,
        )
        mm2_installer.record_store.add_model(fake_config)
        assert mm2_installer.record_store.exists("missing-to-delete")

        # Unregister removes it from DB without touching disk
        mm2_installer.unregister("missing-to-delete")

        with pytest.raises(UnknownModelException):
            mm2_installer.record_store.get_model("missing-to-delete")

    def test_delete_unknown_key_raises(self, mm2_installer: ModelInstallServiceBase) -> None:
        """Deleting a model with an unknown key raises UnknownModelException."""
        with pytest.raises(UnknownModelException):
            mm2_installer.delete("nonexistent-key-12345")

    def test_scan_then_unregister_clears_missing(
        self, mm2_installer: ModelInstallServiceBase, mm2_app_config: InvokeAIAppConfig
    ) -> None:
        """After unregistering all missing models, _scan_for_missing_models returns empty."""
        # Add two models with non-existent paths
        for i in range(2):
            config = TI_File_SD1_Config(
                key=f"missing-bulk-{i}",
                path=f"/nonexistent/bulk_{i}.safetensors",
                name=f"BulkMissing{i}",
                base=BaseModelType.StableDiffusion1,
                type=ModelType.TextualInversion,
                format=ModelFormat.EmbeddingFile,
                hash=f"BULKHASH{i}",
                file_size=1024,
                source=f"test/bulk{i}",
                source_type=ModelSourceType.Path,
            )
            mm2_installer.record_store.add_model(config)

        missing = mm2_installer._scan_for_missing_models()
        assert len(missing) == 2

        # Unregister all missing (simulates bulk delete for missing models)
        for model in missing:
            mm2_installer.unregister(model.key)

        assert len(mm2_installer._scan_for_missing_models()) == 0

    def test_bulk_unregister_does_not_affect_existing_models(
        self,
        mm2_installer: ModelInstallServiceBase,
        embedding_file: Path,
        mm2_app_config: InvokeAIAppConfig,
    ) -> None:
        """Unregistering missing models does not affect models that exist on disk."""
        existing_key = mm2_installer.register_path(embedding_file)

        fake_config = TI_File_SD1_Config(
            key="missing-selective",
            path="/nonexistent/selective.safetensors",
            name="SelectiveMissing",
            base=BaseModelType.StableDiffusion1,
            type=ModelType.TextualInversion,
            format=ModelFormat.EmbeddingFile,
            hash="SELECTIVEHASH",
            file_size=1024,
            source="test/selective",
            source_type=ModelSourceType.Path,
        )
        mm2_installer.record_store.add_model(fake_config)

        # Only unregister the missing one
        missing = mm2_installer._scan_for_missing_models()
        assert len(missing) == 1
        for model in missing:
            mm2_installer.unregister(model.key)

        # Existing model should still be there
        assert mm2_installer.record_store.exists(existing_key)
        assert len(mm2_installer._scan_for_missing_models()) == 0