code/baselines/original/LELA/tests/integration/test_cache_integration.py

"""Integration tests for the caching system."""

import json
import os
import tempfile
import time

import pytest

from lela import Lela
from lela import Document


@pytest.mark.integration
class TestPipelineCacheDirWiring:
    """Test that cache_dir is properly wired through the pipeline."""

    @pytest.fixture
    def config_dict(self, temp_jsonl_kb: str, temp_cache_dir: str) -> dict:
        return {
            "loader": {"name": "text"},
            "ner": {"name": "regex", "params": {"min_len": 3}},
            "candidate_generator": {"name": "fuzzy", "params": {"top_k": 5}},
            "disambiguator": {"name": "first"},
            "knowledge_base": {"name": "jsonl", "params": {"path": temp_jsonl_kb}},
            "cache_dir": temp_cache_dir,
        }

    def test_pipeline_passes_cache_dir_to_kb(
        self, config_dict: dict, temp_cache_dir: str
    ):
        """KB receives cache_dir and creates kb/ cache subdirectory."""
        lela = Lela(config_dict)

        assert lela.kb is not None
        assert hasattr(lela.kb, "source_path")
        assert hasattr(lela.kb, "identity_hash")

        # KB cache should exist
        kb_cache_dir = os.path.join(temp_cache_dir, "kb")
        assert os.path.isdir(kb_cache_dir)
        cache_files = os.listdir(kb_cache_dir)
        assert len(cache_files) == 1
        assert cache_files[0].endswith(".pkl")

    def test_pipeline_kb_cache_used_on_second_init(
        self, config_dict: dict, temp_cache_dir: str
    ):
        """Second pipeline init uses KB cache (produces identical entities)."""
        lela1 = Lela(config_dict)
        entities1 = {e.id: e.title for e in lela1.kb.all_entities()}

        lela2 = Lela(config_dict)
        entities2 = {e.id: e.title for e in lela2.kb.all_entities()}

        assert entities1 == entities2

    def test_pipeline_produces_same_results_with_cache(
        self, config_dict: dict
    ):
        """Pipeline produces identical results on cold and warm runs."""
        lela1 = Lela(config_dict)
        doc = Document(id="test", text="Barack Obama visited New York City yesterday.")
        result1 = lela1.process_document(doc)

        lela2 = Lela(config_dict)
        result2 = lela2.process_document(doc)

        assert result1["text"] == result2["text"]
        assert len(result1["entities"]) == len(result2["entities"])
        for e1, e2 in zip(result1["entities"], result2["entities"]):
            assert e1["text"] == e2["text"]
            assert e1["start"] == e2["start"]
            assert e1["end"] == e2["end"]
            assert e1["entity_id"] == e2["entity_id"]


@pytest.mark.integration
class TestKBCacheInvalidationIntegration:
    """Test KB cache invalidation through the full pipeline."""

    def test_kb_change_invalidates_cache(self, temp_cache_dir: str):
        """Modifying the KB file causes cache rebuild with new data."""
        kb_data = [
            {"id": "Q1", "title": "Barack Obama", "description": "44th US President"},
            {"id": "Q2", "title": "Joe Biden", "description": "46th US President"},
        ]
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".jsonl", delete=False
        ) as f:
            for item in kb_data:
                f.write(json.dumps(item) + "\n")
            kb_path = f.name

        try:
            config_dict = {
                "loader": {"name": "text"},
                "ner": {"name": "regex", "params": {"min_len": 3}},
                "candidate_generator": {"name": "fuzzy", "params": {"top_k": 5}},
                "knowledge_base": {"name": "jsonl", "params": {"path": kb_path}},
                "cache_dir": temp_cache_dir,
            }

            # First run
            lela1 = Lela(config_dict)
            assert len(list(lela1.kb.all_entities())) == 2

            # Modify file
            time.sleep(0.05)
            with open(kb_path, "a") as f:
                f.write(
                    json.dumps({"id": "Q3", "title": "New Entity", "description": "Added"})
                    + "\n"
                )

            # Second run - should see new entity
            lela2 = Lela(config_dict)
            assert len(list(lela2.kb.all_entities())) == 3
            assert lela2.kb.get_entity("Q3") is not None
        finally:
            os.unlink(kb_path)

    def test_clean_cache_rebuilds(self, temp_cache_dir: str, temp_jsonl_kb: str):
        """Deleting the cache directory forces rebuild."""
        import shutil
        from lela.knowledge_bases.jsonl import clear_kb_cache

        config_dict = {
            "loader": {"name": "text"},
            "ner": {"name": "regex", "params": {"min_len": 3}},
            "candidate_generator": {"name": "fuzzy", "params": {"top_k": 5}},
            "knowledge_base": {"name": "jsonl", "params": {"path": temp_jsonl_kb}},
            "cache_dir": temp_cache_dir,
        }

        # Build cache
        lela1 = Lela(config_dict)
        count1 = len(list(lela1.kb.all_entities()))

        # Delete cache contents and clear in-memory cache so rebuild is triggered
        kb_cache = os.path.join(temp_cache_dir, "kb")
        if os.path.isdir(kb_cache):
            shutil.rmtree(kb_cache)
        clear_kb_cache()

        # Rebuild
        lela2 = Lela(config_dict)
        count2 = len(list(lela2.kb.all_entities()))

        assert count1 == count2
        # Cache should be recreated
        assert os.path.isdir(kb_cache)