test_enhanced_crawler.py

#!/usr/bin/env python3
"""
Comprehensive test for the enhanced crawler system.
Tests: frontier queue, entity extraction, LLM logging, PageRank, and all enhancements.
"""

import sys
import time
from pathlib import Path
from datetime import datetime

# Import enhanced components
from database_enhanced import EnhancedCrawlerDatabase
from config import CrawlerConfig


def test_enhanced_database():
    """Test all enhanced database features."""
    print("🧪 Testing Enhanced Crawler Database")
    print("=" * 60)
    
    # Initialize with test database
    test_db_path = Path("test_enhanced.db")
    if test_db_path.exists():
        test_db_path.unlink()
        print(f"🗑️  Removed old test database")
    
    db = EnhancedCrawlerDatabase(test_db_path)
    print(f"✅ Initialized enhanced database at {test_db_path}")
    
    # Test 1: Frontier Queue
    print("\n📋 Test 1: Frontier Queue (Priority-based URL management)")
    print("-" * 60)
    
    test_urls = [
        ("https://example.com/", 0.9, "seed_url", {"source": "manual"}),
        ("https://example.com/about", 0.7, "depth_2", {"depth": 2}),
        ("https://example.com/contact", 0.5, "low_priority", {"depth": 3}),
        ("https://example.com/important", 0.95, "high_value", {"keyword_match": True}),
    ]
    
    for url, score, reason, metadata in test_urls:
        db.add_to_frontier(url, score, reason, metadata)
        print(f"  ➕ Added: {url} (score={score}, reason={reason})")
    
    # Get batch from frontier (should be ordered by score)
    batch = db.get_frontier_batch(limit=2)
    print(f"\n  📦 Retrieved batch (top 2 by score):")
    for item in batch:
        print(f"     - {item['url']} (score={item['score']}, reason={item['reason']})")
    
    # Test 2: Entity Extraction
    print("\n📋 Test 2: Entity Extraction (NER storage)")
    print("-" * 60)
    
    test_page = "https://example.com/about"
    entities = [
        ("PERSON", "John Doe", {"confidence": 0.95}),
        ("ORG", "Example Corporation", {"confidence": 0.88}),
        ("LOCATION", "San Francisco", {"confidence": 0.92}),
        ("DATE", "2025-10-02", {"confidence": 0.99}),
    ]
    
    for kind, label, metadata in entities:
        db.add_entity(test_page, kind, label, metadata)
        print(f"  ➕ Entity: {kind} → {label} (confidence={metadata['confidence']})")
    
    # Retrieve entities
    page_entities = db.get_entities_by_page(test_page)
    print(f"\n  📊 Retrieved {len(page_entities)} entities for {test_page}")
    
    # Get entities by kind
    persons = db.get_entities_by_kind("PERSON")
    print(f"  👤 Found {len(persons)} PERSON entities")
    
    # Test 3: LLM Operation Logging
    print("\n📋 Test 3: LLM Operation Logging (Track LLM usage)")
    print("-" * 60)
    
    llm_ops = [
        ("title_improvement", {"original": "Page Title", "url": test_page}, 
         {"improved": "Enhanced Page Title", "tokens": 150}, 250, True, None),
        ("tag_extraction", {"title": "Enhanced Page Title", "content": "..."}, 
         {"tags": ["web", "example", "test"]}, 180, True, None),
        ("content_summary", {"content": "Long text..."}, 
         {"summary": "Brief summary"}, 300, True, None),
    ]
    
    for op_type, input_data, output_data, duration, success, error in llm_ops:
        db.log_llm_operation(op_type, input_data, output_data, duration, success, error)
        print(f"  📝 Logged: {op_type} ({duration}ms, success={success})")
    
    # Get LLM stats
    stats = db.get_llm_operation_stats()
    print(f"\n  📊 LLM Operation Statistics:")
    if stats:
        total = stats.get('total_ops', 0)
        successful = stats.get('successful_ops', 0)
        success_rate = (successful / total) if total > 0 else 0
        print(f"     Total operations: {total}")
        print(f"     Successful: {successful}")
        print(f"     Failed: {stats.get('failed_ops', 0)}")
        print(f"     Success rate: {success_rate:.1%}")
        print(f"     Average duration: {stats.get('avg_duration_ms', 0):.0f}ms")
        print(f"     Total tokens: {stats.get('total_tokens', 0)}")
    else:
        print("     No LLM operations logged yet")
    
    # Test 4: Crawl Job Tracking
    print("\n📋 Test 4: Crawl Job Tracking (Session management)")
    print("-" * 60)
    
    job_id = db.create_crawl_job("https://example.com", scope_id=None)
    print(f"  🚀 Created crawl job #{job_id}")
    
    # Simulate crawl progress
    time.sleep(0.5)
    job_stats = {
        "pages_crawled": 10,
        "pages_processed": 8,
        "pages_written": 5,
        "duration_seconds": 120
    }
    db.update_crawl_job(job_id, job_stats, status="completed")
    print(f"  ✅ Updated job #{job_id} status to 'completed'")
    print(f"     Stats: {job_stats}")
    
    # Test 5: Enhanced Page Upsert (with slug generation)
    print("\n📝 Test 5: Enhanced Page Upsert (Slug collision fix)")
    print("-" * 60)
    
    import hashlib
    from slugify import slugify
    
    test_pages = [
        {
            "url": "https://example.com/page1",
            "title": "Example Page",
            "content": "Page 1 content",
            "crawl_depth": 1,
        },
        {
            "url": "https://example.com/page2",
            "title": "Example Page",  # Same title as page1
            "content": "Page 2 content",
            "crawl_depth": 2,
        },
    ]
    
    for page_data in test_pages:
        # Generate unique slug with URL hash (same as ContentProcessor does)
        slug = slugify(page_data["title"])
        url_hash = hashlib.sha256(page_data["url"].encode()).hexdigest()[:8]
        page_data["slug"] = f"{slug}-{url_hash}"
        
        result = db.upsert_page_enhanced(page_data["url"], page_data)
        if result:
            print(f"  ✅ Upserted: {page_data['url']}")
            print(f"     Slug: {page_data['slug']}")
        else:
            print(f"  ❌ Failed: {page_data['url']}")
    
    # Test 6: PageRank Computation
    print("\n📝 Test 6: PageRank Computation (Link analysis)")
    print("-" * 60)
    
    # Add links directly to the enhanced database
    # First, ensure the pages exist in the pages table
    link_pages = [
        "https://example.com/",
        "https://example.com/about", 
        "https://example.com/contact",
        "https://example.com/team"
    ]
    
    for url in link_pages:
        slug = slugify(url.rstrip('/') or 'home')
        url_hash = hashlib.sha256(url.encode()).hexdigest()[:8]
        page_data = {
            "url": url,
            "title": f"Page {url.split('/')[-1] or 'Home'}",
            "slug": f"{slug}-{url_hash}",
            "content": f"Content for {url}",
            "crawl_depth": 1,
        }
        db.upsert_page_enhanced(url, page_data)
    
    # Create a simple link structure
    link_structure = {
        "https://example.com/": ["https://example.com/about", "https://example.com/contact"],
        "https://example.com/about": ["https://example.com/", "https://example.com/team"],
        "https://example.com/contact": ["https://example.com/"],
        "https://example.com/team": ["https://example.com/about"],
    }
    
    # Add links directly to the links table (using enhanced schema column names)
    with db._get_connection() as conn:
        cursor = conn.cursor()
        now = datetime.utcnow().isoformat()
        for source, targets in link_structure.items():
            for target in targets:
                cursor.execute("""
                    INSERT OR IGNORE INTO links (src_url, dst_url, rel, anchor_text, first_seen, last_seen)
                    VALUES (?, ?, ?, ?, ?, ?)
                """, (source, target, 'internal', f'Link to {target}', now, now))
        conn.commit()
    
    print(f"  🔗 Created link structure with {len(link_structure)} pages")
    
    # Compute PageRank
    print("  🧮 Computing PageRank...")
    page_ranks = db.compute_page_rank(damping=0.85, max_iter=50)
    
    if page_ranks:
        print(f"\n  📊 PageRank Results (top pages):")
        sorted_ranks = sorted(page_ranks.items(), key=lambda x: x[1], reverse=True)
        for url, rank in sorted_ranks[:5]:
            print(f"     {rank:.4f} - {url}")
    else:
        print(f"  ⚠️  No PageRank results (empty graph)")
    
    # Test 7: Enhanced Statistics
    print("\n📋 Test 7: Enhanced Statistics (Comprehensive metrics)")
    print("-" * 60)
    
    enhanced_stats = db.get_enhanced_statistics()
    
    print(f"  📊 Enhanced Crawler Statistics:")
    print(f"     Total pages: {enhanced_stats.get('total_pages', 0)}")
    print(f"     Processed pages: {enhanced_stats.get('processed_pages', 0)}")
    print(f"     Unique sites: {enhanced_stats.get('unique_sites', 0)}")
    print(f"     Frontier size: {enhanced_stats.get('frontier_size', 0)}")
    print(f"     Average frontier score: {enhanced_stats.get('avg_frontier_score', 0.0)}")
    print(f"     Total entities: {enhanced_stats.get('total_entities', 0)}")
    print(f"     Unique entity kinds: {enhanced_stats.get('unique_entity_kinds', 0)}")
    print(f"     Total links: {enhanced_stats.get('total_links', 0)}")
    print(f"     Internal links: {enhanced_stats.get('internal_links', 0)}")
    print(f"     External links: {enhanced_stats.get('external_links', 0)}")
    
    llm_ops = enhanced_stats.get('llm_operations', {})
    if llm_ops:
        print(f"     LLM operations: {llm_ops.get('total_ops', 0)} total ({llm_ops.get('successful_ops', 0)} successful)")
    
    if enhanced_stats.get('entity_distribution'):
        print(f"\n     Entity Distribution:")
        for kind, count in enhanced_stats['entity_distribution'].items():
            print(f"       {kind}: {count}")
    
    # Final Summary
    print("\n" + "=" * 60)
    print("✅ All Enhanced Database Tests Completed Successfully!")
    print("=" * 60)
    
    print(f"\n📋 Test Summary:")
    print(f"  ✅ Frontier Queue: Working")
    print(f"  ✅ Entity Extraction: Working")
    print(f"  ✅ LLM Logging: Working")
    print(f"  ✅ Crawl Job Tracking: Working")
    print(f"  ✅ Enhanced Page Upsert: Working")
    print(f"  ✅ PageRank Computation: Working")
    print(f"  ✅ Enhanced Statistics: Working")
    
    print(f"\n💾 Test database saved at: {test_db_path}")
    print(f"📊 Database size: {test_db_path.stat().st_size / 1024:.1f} KB")
    
    return True


def test_integration_with_existing_crawler():
    """Test integration with enhanced crawler features."""
    print("\n\n🔗 Testing Enhanced Database API")
    print("=" * 60)
    
    # Create a fresh test database
    test_db_path = Path("test_integration.db")
    if test_db_path.exists():
        test_db_path.unlink()
    
    db = EnhancedCrawlerDatabase(test_db_path)
    
    # Test enhanced API methods
    print("  🔄 Testing enhanced database API...")
    
    import hashlib
    from slugify import slugify
    
    test_url = "https://example.com/test"
    test_data = {
        "title": "Test Page",
        "slug": f"{slugify('Test Page')}-{hashlib.sha256(test_url.encode()).hexdigest()[:8]}",
        "content": "Test content for integration",
        "markdown_content": "# Test Page\n\nTest content for integration",
        "crawl_depth": 1,
        "word_count": 5,
    }
    
    # Test enhanced upsert
    result = db.upsert_page_enhanced(test_url, test_data)
    if result:
        print(f"  ✅ Enhanced upsert_page_enhanced() works")
    else:
        print(f"  ❌ Enhanced upsert failed")
        return False
    
    # Test frontier operations
    db.add_to_frontier(test_url, 0.8, "test", {"test": True})
    frontier = db.get_frontier_batch(limit=1)
    if frontier and len(frontier) > 0:
        print(f"  ✅ Frontier queue operations work")
    else:
        print(f"  ❌ Frontier queue failed")
        return False
    
    # Test entity operations
    db.add_entity(test_url, "PERSON", "Test User", {"confidence": 0.9})
    entities = db.get_entities_by_page(test_url)
    if entities and len(entities) > 0:
        print(f"  ✅ Entity extraction storage works")
    else:
        print(f"  ❌ Entity storage failed")
        return False
    
    # Test LLM logging
    db.log_llm_operation(
        "test_operation",
        {"input": "test"},
        {"output": "result"},
        100,
        True,
        None
    )
    llm_stats = db.get_llm_operation_stats()
    if llm_stats and llm_stats.get('total_ops', 0) > 0:
        print(f"  ✅ LLM operation logging works")
    else:
        print(f"  ❌ LLM logging failed")
        return False
    
    # Get enhanced statistics
    enhanced_stats = db.get_enhanced_statistics()
    
    print(f"\n  📊 Enhanced Statistics:")
    print(f"     Total pages: {enhanced_stats.get('total_pages', 0)}")
    print(f"     Frontier size: {enhanced_stats.get('frontier_size', 0)}")
    print(f"     Total entities: {enhanced_stats.get('total_entities', 0)}")
    llm_ops = enhanced_stats.get('llm_operations', {})
    if llm_ops:
        print(f"     LLM operations: {llm_ops.get('total_ops', 0)} total ({llm_ops.get('successful_ops', 0)} successful)")
    
    print("\n  ✅ Integration test passed - all enhanced features working!")
    
    test_db_path.unlink()
    return True


def main():
    """Run all tests."""
    print("🚀 Enhanced Crawler Comprehensive Test Suite")
    print(f"⏰ Started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("\n")
    
    try:
        # Test 1: Enhanced Database Features
        if not test_enhanced_database():
            print("❌ Enhanced database tests failed")
            return 1
        
        # Test 2: Integration with Existing Crawler
        if not test_integration_with_existing_crawler():
            print("❌ Integration tests failed")
            return 1
        
        # Final success message
        print("\n" + "=" * 60)
        print("🎉 ALL TESTS PASSED SUCCESSFULLY!")
        print("=" * 60)
        print("\n✨ The enhanced crawler is ready for production use.")
        print("\n📝 Next Steps:")
        print("  1. Update main.py to use EnhancedCrawlerDatabase")
        print("  2. Implement enhanced crawler features:")
        print("     - Priority-based crawling with frontier queue")
        print("     - Entity extraction and storage")
        print("     - LLM operation tracking and analytics")
        print("     - PageRank-based page importance")
        print("  3. Test with real crawl on target website")
        print("  4. Build dashboard to visualize enhanced metrics")
        
        return 0
        
    except Exception as e:
        print(f"\n❌ Test suite failed with error: {e}")
        import traceback
        traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())