tests/test_persistence.py

#!/usr/bin/env python3
"""
ChromaDB Persistence Test

Simple test to validate that ChromaDB cache persists across server restarts.
"""

import sys
import os
import time
import tempfile
import shutil
from pathlib import Path

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent.parent))

from cache.chromadb_cache import ChromaDBSearchCache

def test_cache_persistence():
    """Test that cache data persists across restarts"""
    print("Testing ChromaDB Cache Persistence")
    print("=" * 40)
    
    # Create temporary directory for test
    temp_dir = tempfile.mkdtemp(prefix="chromadb_persistence_test_")
    db_path = os.path.join(temp_dir, "test_db")
    results_path = os.path.join(temp_dir, "test_results")
    
    try:
        # Phase 1: Create cache and store data
        print("Phase 1: Creating cache and storing test data...")
        
        cache1 = ChromaDBSearchCache(
            max_size=100,
            default_ttl=3600,
            cache_db_path=db_path,
            cache_results_path=results_path,
            embedding_model="all-MiniLM-L6-v2"
        )
        
        # Store test data
        test_entries = [
            {
                "terms": ["artificial", "intelligence", "machine", "learning"],
                "query": "artificial intelligence machine learning",
                "results": [{"title": "AI Overview", "body": "Introduction to AI", "href": "https://example.com/ai"}]
            },
            {
                "terms": ["python", "programming", "tutorial"],
                "query": "python programming tutorial",
                "results": [{"title": "Python Guide", "body": "Learn Python programming", "href": "https://example.com/python"}]
            },
            {
                "terms": ["web", "development", "javascript"],
                "query": "web development javascript",
                "results": [{"title": "Web Dev", "body": "JavaScript web development", "href": "https://example.com/webdev"}]
            }
        ]
        
        for entry in test_entries:
            cache1.put(entry["terms"], entry["query"], entry["results"])
            print(f"  Stored: {entry['query']}")
        
        stats1 = cache1.get_stats()
        print(f"Cache stats after storing: {stats1['cache_size']} entries")
        
        # Verify data can be retrieved
        print("\nVerifying data retrieval from first cache instance...")
        for entry in test_entries:
            result = cache1.get(entry["terms"])
            if result:
                print(f"  ✅ Found: {entry['query']}")
            else:
                print(f"  ❌ Not found: {entry['query']}")
        
        # Phase 2: Create new cache instance (simulates restart)
        print(f"\nPhase 2: Creating new cache instance (simulating restart)...")
        print("Destroying first cache instance...")
        del cache1  # Remove reference to simulate app restart
        
        time.sleep(1)  # Brief pause
        
        cache2 = ChromaDBSearchCache(
            max_size=100, 
            default_ttl=3600,
            cache_db_path=db_path,
            cache_results_path=results_path,
            embedding_model="all-MiniLM-L6-v2"
        )
        
        stats2 = cache2.get_stats()
        print(f"Cache stats after restart: {stats2['cache_size']} entries")
        
        # Phase 3: Verify persistence
        print(f"\nPhase 3: Verifying data persistence...")
        persisted_count = 0
        
        for entry in test_entries:
            result = cache2.get(entry["terms"])
            if result:
                print(f"  ✅ Persisted: {entry['query']} (age: {time.time() - result.timestamp:.0f}s)")
                persisted_count += 1
            else:
                print(f"  ❌ Lost: {entry['query']}")
        
        # Test semantic similarity after restart
        print(f"\nTesting semantic similarity after restart...")
        
        # Try variations of stored terms
        variations = [
            (["machine", "learning", "artificial", "intelligence"], "Reordered AI terms"),
            (["python", "tutorial"], "Partial Python terms"),
            (["javascript", "web", "development"], "Reordered web terms")
        ]
        
        semantic_hits = 0
        for terms, description in variations:
            result = cache2.get(terms, similarity_threshold=0.7)
            if result:
                print(f"  ✅ Semantic match: {description}")
                semantic_hits += 1
            else:
                print(f"  ❌ No semantic match: {description}")
        
        # Results summary
        print(f"\n" + "="*40)
        print("PERSISTENCE TEST RESULTS")
        print("="*40)
        print(f"Original entries stored: {len(test_entries)}")
        print(f"Entries persisted: {persisted_count}")
        print(f"Persistence rate: {persisted_count / len(test_entries) * 100:.1f}%")
        print(f"Semantic matches: {semantic_hits}/{len(variations)}")
        print(f"Database path: {db_path}")
        print(f"Results path: {results_path}")
        print(f"Database exists: {os.path.exists(db_path)}")
        print(f"Results directory exists: {os.path.exists(results_path)}")
        
        # Check file counts
        if os.path.exists(results_path):
            result_files = [f for f in os.listdir(results_path) if f.endswith('.json')]
            print(f"Result files: {len(result_files)}")
        
        success = persisted_count == len(test_entries)
        print(f"\nPersistence test: {'✅ PASSED' if success else '❌ FAILED'}")
        
        return success
        
    except Exception as e:
        print(f"❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()
        return False
        
    finally:
        # Cleanup
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
        print(f"\nTest cleanup completed")

if __name__ == "__main__":
    success = test_cache_persistence()
    sys.exit(0 if success else 1)