CoolFace
Apppublic

findEthics/Atlas

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
test_persistence.py163 linesDownload Raw Back to tests
1#!/usr/bin/env python32"""3ChromaDB Persistence Test4 5Simple test to validate that ChromaDB cache persists across server restarts.6"""7 8import sys9import os10import time11import tempfile12import shutil13from pathlib import Path14 15# Add parent directory to path for imports16sys.path.append(str(Path(__file__).parent.parent))17 18from cache.chromadb_cache import ChromaDBSearchCache19 20def test_cache_persistence():21    """Test that cache data persists across restarts"""22    print("Testing ChromaDB Cache Persistence")23    print("=" * 40)24    25    # Create temporary directory for test26    temp_dir = tempfile.mkdtemp(prefix="chromadb_persistence_test_")27    db_path = os.path.join(temp_dir, "test_db")28    results_path = os.path.join(temp_dir, "test_results")29    30    try:31        # Phase 1: Create cache and store data32        print("Phase 1: Creating cache and storing test data...")33        34        cache1 = ChromaDBSearchCache(35            max_size=100,36            default_ttl=3600,37            cache_db_path=db_path,38            cache_results_path=results_path,39            embedding_model="all-MiniLM-L6-v2"40        )41        42        # Store test data43        test_entries = [44            {45                "terms": ["artificial", "intelligence", "machine", "learning"],46                "query": "artificial intelligence machine learning",47                "results": [{"title": "AI Overview", "body": "Introduction to AI", "href": "https://example.com/ai"}]48            },49            {50                "terms": ["python", "programming", "tutorial"],51                "query": "python programming tutorial",52                "results": [{"title": "Python Guide", "body": "Learn Python programming", "href": "https://example.com/python"}]53            },54            {55                "terms": ["web", "development", "javascript"],56                "query": "web development javascript",57                "results": [{"title": "Web Dev", "body": "JavaScript web development", "href": "https://example.com/webdev"}]58            }59        ]60        61        for entry in test_entries:62            cache1.put(entry["terms"], entry["query"], entry["results"])63            print(f"  Stored: {entry['query']}")64        65        stats1 = cache1.get_stats()66        print(f"Cache stats after storing: {stats1['cache_size']} entries")67        68        # Verify data can be retrieved69        print("\nVerifying data retrieval from first cache instance...")70        for entry in test_entries:71            result = cache1.get(entry["terms"])72            if result:73                print(f"  ✅ Found: {entry['query']}")74            else:75                print(f"  ❌ Not found: {entry['query']}")76        77        # Phase 2: Create new cache instance (simulates restart)78        print(f"\nPhase 2: Creating new cache instance (simulating restart)...")79        print("Destroying first cache instance...")80        del cache1  # Remove reference to simulate app restart81        82        time.sleep(1)  # Brief pause83        84        cache2 = ChromaDBSearchCache(85            max_size=100, 86            default_ttl=3600,87            cache_db_path=db_path,88            cache_results_path=results_path,89            embedding_model="all-MiniLM-L6-v2"90        )91        92        stats2 = cache2.get_stats()93        print(f"Cache stats after restart: {stats2['cache_size']} entries")94        95        # Phase 3: Verify persistence96        print(f"\nPhase 3: Verifying data persistence...")97        persisted_count = 098        99        for entry in test_entries:100            result = cache2.get(entry["terms"])101            if result:102                print(f"  ✅ Persisted: {entry['query']} (age: {time.time() - result.timestamp:.0f}s)")103                persisted_count += 1104            else:105                print(f"  ❌ Lost: {entry['query']}")106        107        # Test semantic similarity after restart108        print(f"\nTesting semantic similarity after restart...")109        110        # Try variations of stored terms111        variations = [112            (["machine", "learning", "artificial", "intelligence"], "Reordered AI terms"),113            (["python", "tutorial"], "Partial Python terms"),114            (["javascript", "web", "development"], "Reordered web terms")115        ]116        117        semantic_hits = 0118        for terms, description in variations:119            result = cache2.get(terms, similarity_threshold=0.7)120            if result:121                print(f"  ✅ Semantic match: {description}")122                semantic_hits += 1123            else:124                print(f"  ❌ No semantic match: {description}")125        126        # Results summary127        print(f"\n" + "="*40)128        print("PERSISTENCE TEST RESULTS")129        print("="*40)130        print(f"Original entries stored: {len(test_entries)}")131        print(f"Entries persisted: {persisted_count}")132        print(f"Persistence rate: {persisted_count / len(test_entries) * 100:.1f}%")133        print(f"Semantic matches: {semantic_hits}/{len(variations)}")134        print(f"Database path: {db_path}")135        print(f"Results path: {results_path}")136        print(f"Database exists: {os.path.exists(db_path)}")137        print(f"Results directory exists: {os.path.exists(results_path)}")138        139        # Check file counts140        if os.path.exists(results_path):141            result_files = [f for f in os.listdir(results_path) if f.endswith('.json')]142            print(f"Result files: {len(result_files)}")143        144        success = persisted_count == len(test_entries)145        print(f"\nPersistence test: {'✅ PASSED' if success else '❌ FAILED'}")146        147        return success148        149    except Exception as e:150        print(f"❌ Test failed with error: {e}")151        import traceback152        traceback.print_exc()153        return False154        155    finally:156        # Cleanup157        if os.path.exists(temp_dir):158            shutil.rmtree(temp_dir)159        print(f"\nTest cleanup completed")160 161if __name__ == "__main__":162    success = test_cache_persistence()163    sys.exit(0 if success else 1)