File size: 5,445 Bytes
ae11455 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """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)
|