ComfyUI-Lora-Manager/tests/test_recipe_fts_index_validation.py

"""Tests for RecipeFTSIndex validation methods."""

import os
import tempfile
from typing import Any, Dict, List

import pytest

from py.services.recipe_fts_index import RecipeFTSIndex


@pytest.fixture
def temp_db_path():
    """Create a temporary database path."""
    with tempfile.NamedTemporaryFile(suffix=".sqlite", delete=False) as f:
        path = f.name
    yield path
    # Cleanup
    if os.path.exists(path):
        os.unlink(path)
    for suffix in ["-wal", "-shm"]:
        wal_path = path + suffix
        if os.path.exists(wal_path):
            os.unlink(wal_path)


@pytest.fixture
def sample_recipes() -> List[Dict[str, Any]]:
    """Create sample recipe data for FTS indexing."""
    return [
        {
            "id": "recipe-001",
            "title": "Anime Character Portrait",
            "tags": ["anime", "portrait", "character"],
            "loras": [
                {"file_name": "anime_style", "modelName": "Anime Style LoRA"},
                {"file_name": "character_v2", "modelName": "Character Design V2"},
            ],
            "gen_params": {
                "prompt": "masterpiece, best quality, 1girl",
                "negative_prompt": "bad quality, worst quality",
            },
        },
        {
            "id": "recipe-002",
            "title": "Landscape Photography",
            "tags": ["landscape", "photography", "nature"],
            "loras": [
                {"file_name": "landscape_lora", "modelName": "Landscape Enhancement"},
            ],
            "gen_params": {
                "prompt": "beautiful landscape, mountains, sunset",
                "negative_prompt": "ugly, blurry",
            },
        },
        {
            "id": "recipe-003",
            "title": "Fantasy Art Scene",
            "tags": ["fantasy", "art"],
            "loras": [],
            "gen_params": {
                "prompt": "fantasy world, dragons, magic",
            },
        },
    ]


class TestFTSIndexValidation:
    """Tests for FTS index validation methods."""

    def test_validate_index_empty_returns_false(self, temp_db_path):
        """Test that validation fails on empty index."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.initialize()

        # Empty index should not validate against non-empty recipe set
        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is False

    def test_validate_index_count_mismatch(self, temp_db_path, sample_recipes):
        """Test validation fails when counts don't match."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)

        # Validate with wrong count
        result = fts.validate_index(5, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is False

    def test_validate_index_id_mismatch(self, temp_db_path, sample_recipes):
        """Test validation fails when IDs don't match."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)

        # Validate with wrong IDs
        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-999"})
        assert result is False

    def test_validate_index_success(self, temp_db_path, sample_recipes):
        """Test successful validation."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)

        # Validate with correct count and IDs
        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is True

    def test_get_indexed_recipe_ids(self, temp_db_path, sample_recipes):
        """Test getting indexed recipe IDs."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)

        ids = fts.get_indexed_recipe_ids()
        assert ids == {"recipe-001", "recipe-002", "recipe-003"}

    def test_get_indexed_recipe_ids_empty(self, temp_db_path):
        """Test getting IDs from empty index."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.initialize()

        ids = fts.get_indexed_recipe_ids()
        assert ids == set()

    def test_validate_after_add_recipe(self, temp_db_path, sample_recipes):
        """Test validation after adding a recipe."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes[:2])  # Only first 2

        # Validation should fail with all 3 IDs
        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is False

        # Add third recipe
        fts.add_recipe(sample_recipes[2])

        # Now validation should pass
        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is True

    def test_validate_after_remove_recipe(self, temp_db_path, sample_recipes):
        """Test validation after removing a recipe."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)

        # Remove a recipe
        fts.remove_recipe("recipe-002")

        # Validation should fail with original 3 IDs
        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is False

        # Validation should pass with 2 remaining IDs
        result = fts.validate_index(2, {"recipe-001", "recipe-003"})
        assert result is True

    def test_validate_index_uninitialized(self, temp_db_path):
        """Test validation on uninitialized index."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        # Don't call initialize

        # Should initialize automatically and return False for non-empty set
        result = fts.validate_index(1, {"recipe-001"})
        assert result is False

    def test_indexed_count_after_clear(self, temp_db_path, sample_recipes):
        """Test count after clearing index."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)
        assert fts.get_indexed_count() == 3

        fts.clear()
        assert fts.get_indexed_count() == 0

    def test_search_still_works_after_validation(self, temp_db_path, sample_recipes):
        """Test that search works correctly after validation."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)

        # Validate (which checks state)
        fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"})

        # Search should still work
        results = fts.search("anime")
        assert "recipe-001" in results

    @staticmethod
    def _forbid_content_scan(monkeypatch):
        """Fail the test if validation falls back to scanning the FTS content table."""

        def _raise(*args, **kwargs):
            raise AssertionError("validate_index scanned the FTS content table")

        monkeypatch.setattr(RecipeFTSIndex, "get_indexed_count", _raise)
        monkeypatch.setattr(RecipeFTSIndex, "get_indexed_recipe_ids", _raise)

    def test_validate_index_with_metadata_does_not_scan(
        self, temp_db_path, sample_recipes, monkeypatch
    ):
        """Validation must rely on stored metadata, not content-table scans."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)

        self._forbid_content_scan(monkeypatch)

        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is True

        # Mismatches are also detected from metadata alone
        result = fts.validate_index(4, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is False
        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-999"})
        assert result is False

    def test_validate_index_legacy_fallback_writes_metadata(
        self, temp_db_path, sample_recipes, monkeypatch
    ):
        """Indexes built by older versions validate via a one-time scan, then metadata."""
        import sqlite3

        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)

        # Simulate a legacy index: drop the fingerprint metadata
        conn = sqlite3.connect(temp_db_path)
        try:
            conn.execute(
                "DELETE FROM fts_metadata WHERE key = ?",
                (RecipeFTSIndex._FINGERPRINT_METADATA_KEY,)
            )
            conn.commit()
        finally:
            conn.close()

        # Fallback validation scans once and succeeds
        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is True

        # Metadata was written, so the next validation needs no scan
        self._forbid_content_scan(monkeypatch)
        result = fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"})
        assert result is True

    def test_validate_index_metadata_tracks_incremental_mutations(
        self, temp_db_path, sample_recipes, monkeypatch
    ):
        """add_recipe/remove_recipe keep the validation metadata in sync."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes[:2])

        fts.add_recipe(sample_recipes[2])
        fts.remove_recipe("recipe-001")

        self._forbid_content_scan(monkeypatch)

        result = fts.validate_index(2, {"recipe-002", "recipe-003"})
        assert result is True

    def test_validate_index_after_clear_uses_metadata(
        self, temp_db_path, sample_recipes, monkeypatch
    ):
        """clear() resets the validation metadata to the empty index state."""
        fts = RecipeFTSIndex(db_path=temp_db_path)
        fts.build_index(sample_recipes)
        fts.clear()

        self._forbid_content_scan(monkeypatch)

        assert fts.validate_index(0, set()) is True
        assert fts.validate_index(3, {"recipe-001", "recipe-002", "recipe-003"}) is False