mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
e4a41faa33
- Fix E402: Move module-level imports to top of file - Fix F401: Add noqa for availability check imports - Fix F402: Rename loop variables shadowing imports - Fix E722: Replace bare except with except Exception - Fix B904: Add exception chaining (from e) - Fix F811: Remove duplicate imports - Fix B027: Add noqa for empty close() method - Fix E741: Rename ambiguous variable l -> label - Fix I001: Import sorting issues - Apply ruff format to all 106 files All 902 tests pass.
617 lines
19 KiB
Python
617 lines
19 KiB
Python
"""Tests for Compress-Cache-Retrieve (CCR) architecture.
|
|
|
|
These tests verify that:
|
|
1. CompressionStore correctly caches compressed content
|
|
2. SmartCrusher integrates with CompressionStore
|
|
3. Retrieval works correctly (full and search)
|
|
4. Feedback tracking works
|
|
5. TTL expiration works
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from headroom.cache.compression_store import (
|
|
CompressionStore,
|
|
get_compression_store,
|
|
reset_compression_store,
|
|
)
|
|
from headroom.config import CCRConfig
|
|
from headroom.transforms.smart_crusher import (
|
|
SmartCrusherConfig,
|
|
smart_crush_tool_output,
|
|
)
|
|
|
|
|
|
class TestCompressionStore:
|
|
"""Test CompressionStore functionality."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_store(self):
|
|
"""Reset global store before each test."""
|
|
reset_compression_store()
|
|
yield
|
|
reset_compression_store()
|
|
|
|
def test_store_and_retrieve(self):
|
|
"""Basic store and retrieve flow."""
|
|
store = CompressionStore()
|
|
|
|
original = json.dumps([{"id": i} for i in range(100)])
|
|
compressed = json.dumps([{"id": i} for i in range(10)])
|
|
|
|
hash_key = store.store(
|
|
original=original,
|
|
compressed=compressed,
|
|
original_tokens=1000,
|
|
compressed_tokens=100,
|
|
original_item_count=100,
|
|
compressed_item_count=10,
|
|
)
|
|
|
|
assert (
|
|
len(hash_key) == 24
|
|
) # SHA256 truncated to 24 chars (96 bits for collision resistance)
|
|
|
|
entry = store.retrieve(hash_key)
|
|
assert entry is not None
|
|
assert entry.original_content == original
|
|
assert entry.compressed_content == compressed
|
|
assert entry.original_tokens == 1000
|
|
assert entry.compressed_tokens == 100
|
|
|
|
def test_retrieve_nonexistent(self):
|
|
"""Retrieve returns None for nonexistent hash."""
|
|
store = CompressionStore()
|
|
entry = store.retrieve("nonexistent1234")
|
|
assert entry is None
|
|
|
|
def test_ttl_expiration(self):
|
|
"""Entries expire after TTL."""
|
|
store = CompressionStore(default_ttl=1) # 1 second TTL
|
|
|
|
hash_key = store.store(
|
|
original="[1,2,3]",
|
|
compressed="[1]",
|
|
ttl=1,
|
|
)
|
|
|
|
# Should exist immediately
|
|
assert store.exists(hash_key)
|
|
|
|
# Wait for expiration
|
|
time.sleep(1.1)
|
|
|
|
# Should be expired
|
|
assert not store.exists(hash_key)
|
|
entry = store.retrieve(hash_key)
|
|
assert entry is None
|
|
|
|
def test_eviction_at_capacity(self):
|
|
"""Oldest entries evicted when at capacity."""
|
|
store = CompressionStore(max_entries=3)
|
|
|
|
hashes = []
|
|
for i in range(5):
|
|
h = store.store(
|
|
original=f"original_{i}",
|
|
compressed=f"compressed_{i}",
|
|
)
|
|
hashes.append(h)
|
|
time.sleep(0.01) # Ensure different timestamps
|
|
|
|
# Only last 3 should exist
|
|
assert not store.exists(hashes[0])
|
|
assert not store.exists(hashes[1])
|
|
assert store.exists(hashes[2])
|
|
assert store.exists(hashes[3])
|
|
assert store.exists(hashes[4])
|
|
|
|
def test_search_with_bm25(self):
|
|
"""Search within cached content using BM25."""
|
|
store = CompressionStore()
|
|
|
|
items = [
|
|
{"id": 1, "content": "Python programming language"},
|
|
{"id": 2, "content": "JavaScript web development"},
|
|
{"id": 3, "content": "Python data science pandas"},
|
|
{"id": 4, "content": "Java enterprise applications"},
|
|
{"id": 5, "content": "Python machine learning tensorflow"},
|
|
]
|
|
|
|
hash_key = store.store(
|
|
original=json.dumps(items),
|
|
compressed=json.dumps(items[:2]),
|
|
original_item_count=5,
|
|
compressed_item_count=2,
|
|
)
|
|
|
|
# Search for Python items
|
|
results = store.search(hash_key, "Python programming")
|
|
|
|
assert len(results) >= 1
|
|
# Should prioritize Python items
|
|
result_ids = [r["id"] for r in results]
|
|
assert 1 in result_ids # "Python programming language"
|
|
|
|
def test_retrieval_tracking(self):
|
|
"""Retrieval events are tracked for feedback."""
|
|
store = CompressionStore(enable_feedback=True)
|
|
|
|
hash_key = store.store(
|
|
original="[1,2,3]",
|
|
compressed="[1]",
|
|
tool_name="test_tool",
|
|
)
|
|
|
|
# Retrieve multiple times
|
|
store.retrieve(hash_key)
|
|
store.retrieve(hash_key, query="test query")
|
|
store.search(hash_key, "another query")
|
|
|
|
events = store.get_retrieval_events(limit=10)
|
|
assert len(events) >= 2
|
|
|
|
# Check event details
|
|
assert any(e.retrieval_type == "full" for e in events)
|
|
assert any(e.retrieval_type == "search" for e in events)
|
|
|
|
def test_access_tracking_on_entry(self):
|
|
"""Entry tracks access count and queries."""
|
|
store = CompressionStore()
|
|
|
|
hash_key = store.store(
|
|
original=json.dumps([{"id": i} for i in range(10)]),
|
|
compressed="[]",
|
|
)
|
|
|
|
# Access multiple times with queries
|
|
store.retrieve(hash_key, query="first query")
|
|
store.retrieve(hash_key, query="second query")
|
|
store.retrieve(hash_key, query="first query") # Duplicate
|
|
|
|
entry = store.retrieve(hash_key)
|
|
assert entry.retrieval_count >= 3
|
|
assert "first query" in entry.search_queries
|
|
assert "second query" in entry.search_queries
|
|
|
|
def test_stats(self):
|
|
"""Store statistics are accurate."""
|
|
store = CompressionStore()
|
|
|
|
store.store(
|
|
original="x" * 100,
|
|
compressed="x" * 10,
|
|
original_tokens=100,
|
|
compressed_tokens=10,
|
|
)
|
|
store.store(
|
|
original="y" * 200,
|
|
compressed="y" * 20,
|
|
original_tokens=200,
|
|
compressed_tokens=20,
|
|
)
|
|
|
|
stats = store.get_stats()
|
|
assert stats["entry_count"] == 2
|
|
assert stats["total_original_tokens"] == 300
|
|
assert stats["total_compressed_tokens"] == 30
|
|
|
|
def test_global_store_singleton(self):
|
|
"""Global store uses singleton pattern."""
|
|
reset_compression_store()
|
|
|
|
store1 = get_compression_store()
|
|
store2 = get_compression_store()
|
|
|
|
assert store1 is store2
|
|
|
|
def test_thread_safety(self):
|
|
"""Store is thread-safe."""
|
|
import threading
|
|
|
|
store = CompressionStore()
|
|
hashes = []
|
|
lock = threading.Lock()
|
|
|
|
def store_item(i):
|
|
h = store.store(
|
|
original=f"original_{i}",
|
|
compressed=f"compressed_{i}",
|
|
)
|
|
with lock:
|
|
hashes.append(h)
|
|
|
|
threads = [threading.Thread(target=store_item, args=(i,)) for i in range(10)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert len(hashes) == 10
|
|
for h in hashes:
|
|
assert store.exists(h)
|
|
|
|
|
|
class TestSmartCrusherCCRIntegration:
|
|
"""Test SmartCrusher integration with CCR."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_store(self):
|
|
"""Reset global store before each test."""
|
|
reset_compression_store()
|
|
yield
|
|
reset_compression_store()
|
|
|
|
def test_compression_caches_original(self):
|
|
"""SmartCrusher caches original content when compressing."""
|
|
items = [{"id": i, "score": 100 - i, "data": f"item_{i}"} for i in range(100)]
|
|
content = json.dumps(items)
|
|
|
|
config = SmartCrusherConfig(max_items_after_crush=15)
|
|
ccr_config = CCRConfig(
|
|
enabled=True,
|
|
inject_retrieval_marker=False, # Don't add marker for this test
|
|
min_items_to_cache=10,
|
|
)
|
|
|
|
compressed_str, was_modified, _ = smart_crush_tool_output(content, config, ccr_config)
|
|
|
|
assert was_modified
|
|
|
|
# Check that content was cached
|
|
store = get_compression_store()
|
|
stats = store.get_stats()
|
|
assert stats["entry_count"] >= 1
|
|
|
|
def test_retrieval_marker_injected(self):
|
|
"""CCR marker is injected when configured."""
|
|
items = [{"id": i, "score": 100 - i, "data": f"item_{i}"} for i in range(100)]
|
|
content = json.dumps(items)
|
|
|
|
config = SmartCrusherConfig(max_items_after_crush=15)
|
|
ccr_config = CCRConfig(
|
|
enabled=True,
|
|
inject_retrieval_marker=True,
|
|
min_items_to_cache=10,
|
|
)
|
|
|
|
compressed_str, was_modified, _ = smart_crush_tool_output(content, config, ccr_config)
|
|
|
|
assert was_modified
|
|
# Marker should be present
|
|
assert "items compressed" in compressed_str or "hash=" in compressed_str
|
|
|
|
def test_small_arrays_not_cached(self):
|
|
"""Arrays smaller than min_items_to_cache are not cached."""
|
|
items = [{"id": i} for i in range(15)]
|
|
content = json.dumps(items)
|
|
|
|
config = SmartCrusherConfig(max_items_after_crush=10)
|
|
ccr_config = CCRConfig(
|
|
enabled=True,
|
|
min_items_to_cache=50, # Require 50+ items
|
|
)
|
|
|
|
smart_crush_tool_output(content, config, ccr_config)
|
|
|
|
store = get_compression_store()
|
|
stats = store.get_stats()
|
|
# Should not cache because original has < 50 items
|
|
assert stats["entry_count"] == 0
|
|
|
|
def test_uncrushed_data_not_cached(self):
|
|
"""Data that doesn't get crushed is not cached."""
|
|
# DB results with unique IDs - shouldn't be crushed
|
|
items = [{"id": i, "name": f"User {i}", "email": f"user{i}@test.com"} for i in range(30)]
|
|
content = json.dumps(items)
|
|
|
|
config = SmartCrusherConfig(max_items_after_crush=10)
|
|
ccr_config = CCRConfig(enabled=True, min_items_to_cache=10)
|
|
|
|
compressed_str, was_modified, _ = smart_crush_tool_output(content, config, ccr_config)
|
|
|
|
# If not modified, shouldn't be cached
|
|
if not was_modified:
|
|
store = get_compression_store()
|
|
stats = store.get_stats()
|
|
assert stats["entry_count"] == 0
|
|
|
|
def test_can_retrieve_after_compression(self):
|
|
"""Can retrieve original content after compression."""
|
|
items = [
|
|
{"id": i, "score": 100 - i, "content": f"Document about topic {i}"} for i in range(100)
|
|
]
|
|
content = json.dumps(items)
|
|
|
|
config = SmartCrusherConfig(max_items_after_crush=15)
|
|
ccr_config = CCRConfig(
|
|
enabled=True,
|
|
inject_retrieval_marker=True,
|
|
min_items_to_cache=10,
|
|
)
|
|
|
|
compressed_str, was_modified, _ = smart_crush_tool_output(content, config, ccr_config)
|
|
|
|
assert was_modified
|
|
|
|
# Extract hash from marker
|
|
# Marker format: [100 items compressed to 15. Retrieve more: hash=abc123...]
|
|
import re
|
|
|
|
match = re.search(r"hash=([a-f0-9]+)", compressed_str)
|
|
assert match is not None, f"No hash found in: {compressed_str}"
|
|
|
|
hash_key = match.group(1)
|
|
|
|
# Retrieve original
|
|
store = get_compression_store()
|
|
entry = store.retrieve(hash_key)
|
|
|
|
assert entry is not None
|
|
original_items = json.loads(entry.original_content)
|
|
assert len(original_items) == 100
|
|
|
|
def test_search_after_compression(self):
|
|
"""Can search within original content after compression."""
|
|
items = [
|
|
{"id": 1, "content": "Authentication error: invalid token"},
|
|
{"id": 2, "content": "Database connection successful"},
|
|
{"id": 3, "content": "User login completed"},
|
|
] + [{"id": i, "content": f"Generic log entry {i}"} for i in range(4, 104)]
|
|
content = json.dumps(items)
|
|
|
|
config = SmartCrusherConfig(max_items_after_crush=15)
|
|
ccr_config = CCRConfig(
|
|
enabled=True,
|
|
inject_retrieval_marker=True,
|
|
min_items_to_cache=10,
|
|
)
|
|
|
|
compressed_str, was_modified, _ = smart_crush_tool_output(content, config, ccr_config)
|
|
|
|
assert was_modified
|
|
|
|
# Extract hash
|
|
import re
|
|
|
|
match = re.search(r"hash=([a-f0-9]+)", compressed_str)
|
|
hash_key = match.group(1)
|
|
|
|
# Search for authentication items
|
|
store = get_compression_store()
|
|
results = store.search(hash_key, "authentication error token")
|
|
|
|
# Should find the authentication error item
|
|
assert len(results) >= 1
|
|
assert any("Authentication" in r.get("content", "") for r in results)
|
|
|
|
|
|
class TestCCRConfig:
|
|
"""Test CCR configuration options."""
|
|
|
|
def test_default_config(self):
|
|
"""Default CCR config values."""
|
|
config = CCRConfig()
|
|
assert config.enabled is True
|
|
assert config.store_max_entries == 1000
|
|
assert config.store_ttl_seconds == 300
|
|
assert config.inject_retrieval_marker is True
|
|
assert config.feedback_enabled is True
|
|
assert config.min_items_to_cache == 20
|
|
|
|
def test_custom_marker_template(self):
|
|
"""Custom marker template is used."""
|
|
items = [{"id": i, "score": 100 - i} for i in range(100)]
|
|
content = json.dumps(items)
|
|
|
|
config = SmartCrusherConfig(max_items_after_crush=15)
|
|
ccr_config = CCRConfig(
|
|
enabled=True,
|
|
inject_retrieval_marker=True,
|
|
min_items_to_cache=10,
|
|
marker_template="\n[CUSTOM: {original_count} -> {compressed_count}, key={hash}]",
|
|
)
|
|
|
|
compressed_str, was_modified, _ = smart_crush_tool_output(content, config, ccr_config)
|
|
|
|
if was_modified:
|
|
assert "CUSTOM:" in compressed_str or "key=" in compressed_str
|
|
|
|
|
|
class TestCCRFeedbackLoop:
|
|
"""Test CCR feedback tracking for learning."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_store(self):
|
|
"""Reset global store before each test."""
|
|
reset_compression_store()
|
|
yield
|
|
reset_compression_store()
|
|
|
|
def test_retrieval_events_logged(self):
|
|
"""Retrieval events are logged for feedback."""
|
|
store = CompressionStore(enable_feedback=True)
|
|
|
|
items = [{"id": i, "data": f"item_{i}"} for i in range(50)]
|
|
hash_key = store.store(
|
|
original=json.dumps(items),
|
|
compressed=json.dumps(items[:10]),
|
|
original_item_count=50,
|
|
compressed_item_count=10,
|
|
tool_name="search_api",
|
|
)
|
|
|
|
# Simulate retrievals
|
|
store.retrieve(hash_key)
|
|
store.search(hash_key, "specific query")
|
|
store.search(hash_key, "another query")
|
|
|
|
events = store.get_retrieval_events(limit=10)
|
|
|
|
# Should have logged all retrievals
|
|
assert len(events) >= 3
|
|
|
|
# Check event types
|
|
full_events = [e for e in events if e.retrieval_type == "full"]
|
|
search_events = [e for e in events if e.retrieval_type == "search"]
|
|
|
|
assert len(full_events) >= 1
|
|
assert len(search_events) >= 2
|
|
|
|
def test_tool_name_in_events(self):
|
|
"""Tool name is preserved in retrieval events."""
|
|
store = CompressionStore(enable_feedback=True)
|
|
|
|
hash_key = store.store(
|
|
original="[1,2,3]",
|
|
compressed="[1]",
|
|
tool_name="github_search",
|
|
)
|
|
|
|
store.retrieve(hash_key)
|
|
|
|
events = store.get_retrieval_events(tool_name="github_search")
|
|
assert len(events) >= 1
|
|
assert all(e.tool_name == "github_search" for e in events)
|
|
|
|
def test_event_filtering_by_tool(self):
|
|
"""Events can be filtered by tool name."""
|
|
store = CompressionStore(enable_feedback=True)
|
|
|
|
hash1 = store.store(
|
|
original="[1]",
|
|
compressed="[1]",
|
|
tool_name="tool_a",
|
|
)
|
|
hash2 = store.store(
|
|
original="[2]",
|
|
compressed="[2]",
|
|
tool_name="tool_b",
|
|
)
|
|
|
|
store.retrieve(hash1)
|
|
store.retrieve(hash1)
|
|
store.retrieve(hash2)
|
|
|
|
tool_a_events = store.get_retrieval_events(tool_name="tool_a")
|
|
tool_b_events = store.get_retrieval_events(tool_name="tool_b")
|
|
|
|
assert len(tool_a_events) == 2
|
|
assert len(tool_b_events) == 1
|
|
|
|
|
|
class TestCCREdgeCases:
|
|
"""Test edge cases and error handling."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_store(self):
|
|
"""Reset global store before each test."""
|
|
reset_compression_store()
|
|
yield
|
|
reset_compression_store()
|
|
|
|
def test_search_expired_entry(self):
|
|
"""Search on expired entry returns empty."""
|
|
store = CompressionStore(default_ttl=1)
|
|
|
|
hash_key = store.store(
|
|
original=json.dumps([{"id": 1}]),
|
|
compressed="[]",
|
|
)
|
|
|
|
time.sleep(1.1)
|
|
|
|
results = store.search(hash_key, "query")
|
|
assert results == []
|
|
|
|
def test_search_invalid_json(self):
|
|
"""Search handles invalid JSON gracefully."""
|
|
store = CompressionStore()
|
|
|
|
hash_key = store.store(
|
|
original="not valid json",
|
|
compressed="[]",
|
|
)
|
|
|
|
results = store.search(hash_key, "query")
|
|
assert results == []
|
|
|
|
def test_search_non_array(self):
|
|
"""Search handles non-array content gracefully."""
|
|
store = CompressionStore()
|
|
|
|
hash_key = store.store(
|
|
original=json.dumps({"key": "value"}),
|
|
compressed="{}",
|
|
)
|
|
|
|
results = store.search(hash_key, "query")
|
|
assert results == []
|
|
|
|
def test_empty_query_search(self):
|
|
"""Search with empty query returns empty or all."""
|
|
store = CompressionStore()
|
|
|
|
items = [{"id": i} for i in range(10)]
|
|
hash_key = store.store(
|
|
original=json.dumps(items),
|
|
compressed="[]",
|
|
)
|
|
|
|
# Empty query should return something (BM25 handles this)
|
|
results = store.search(hash_key, "")
|
|
# Behavior depends on BM25 implementation
|
|
assert isinstance(results, list)
|
|
|
|
def test_ccr_disabled_no_caching(self):
|
|
"""When CCR disabled, no caching occurs."""
|
|
reset_compression_store()
|
|
|
|
items = [{"id": i, "score": 100 - i} for i in range(100)]
|
|
content = json.dumps(items)
|
|
|
|
config = SmartCrusherConfig(max_items_after_crush=15)
|
|
ccr_config = CCRConfig(enabled=False) # Disabled
|
|
|
|
smart_crush_tool_output(content, config, ccr_config)
|
|
|
|
store = get_compression_store()
|
|
stats = store.get_stats()
|
|
assert stats["entry_count"] == 0
|
|
|
|
def test_concurrent_store_and_retrieve(self):
|
|
"""Concurrent operations don't corrupt data."""
|
|
import threading
|
|
|
|
store = CompressionStore()
|
|
errors = []
|
|
|
|
def store_and_retrieve(i):
|
|
try:
|
|
items = [{"id": j, "batch": i} for j in range(10)]
|
|
hash_key = store.store(
|
|
original=json.dumps(items),
|
|
compressed="[]",
|
|
tool_name=f"tool_{i}",
|
|
)
|
|
|
|
# Immediately retrieve
|
|
entry = store.retrieve(hash_key)
|
|
if entry is None:
|
|
errors.append(f"Entry {i} not found after store")
|
|
elif f'"batch": {i}' not in entry.original_content:
|
|
errors.append(f"Entry {i} has wrong content")
|
|
except Exception as e:
|
|
errors.append(str(e))
|
|
|
|
threads = [threading.Thread(target=store_and_retrieve, args=(i,)) for i in range(20)]
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert errors == [], f"Errors during concurrent operations: {errors}"
|