test_complete_hf_deployment.py

#!/usr/bin/env python3
"""
Test script to validate complete HF deployment functionality
Tests file ingestion, RAG pipeline, and query responses
"""

import os
import sys

sys.path.append("src")


def test_file_ingestion():
    """Test that synthetic_policies files are being processed"""
    print("🔍 Testing file ingestion...")

    # Check if synthetic_policies directory exists
    corpus_dir = "synthetic_policies"
    if not os.path.exists(corpus_dir):
        print(f"❌ Corpus directory not found: {corpus_dir}")
        return False

    # Count files in directory
    policy_files = []
    for file in os.listdir(corpus_dir):
        if file.endswith((".md", ".txt")):
            policy_files.append(file)

    print(f"✅ Found {len(policy_files)} policy files in {corpus_dir}")
    return len(policy_files) > 0


def test_rag_pipeline_initialization():
    """Test RAG pipeline initialization with ingestion"""
    print("🔍 Testing RAG pipeline initialization...")

    try:
        # Set environment to enable HF services
        os.environ["OPENROUTER_API_KEY"] = "test"

        from routes.main_routes import get_rag_pipeline

        pipeline = get_rag_pipeline()

        if pipeline is None:
            print("❌ RAG pipeline not initialized")
            return False

        print("✅ RAG pipeline initialized")

        # Check search service
        if not hasattr(pipeline, "search_service") or pipeline.search_service is None:
            print("❌ Search service not available")
            return False

        print("✅ Search service available")

        # Check vector store
        if hasattr(pipeline.search_service, "vector_db"):
            try:
                count = pipeline.search_service.vector_db.get_count()
                print(f"✅ Vector store has {count} embeddings")
                return count > 0
            except Exception as e:
                print(f"⚠️ Could not get embedding count: {e}")
                return False
        else:
            print("❌ Vector database not accessible")
            return False

    except Exception as e:
        print(f"❌ RAG pipeline initialization failed: {e}")
        return False


def test_query_response():
    """Test that queries return proper responses"""
    print("🔍 Testing query response functionality...")

    try:
        # Set environment to enable HF services
        os.environ["OPENROUTER_API_KEY"] = "test"

        from routes.main_routes import get_rag_pipeline

        pipeline = get_rag_pipeline()

        if pipeline is None or pipeline.search_service is None:
            print("❌ RAG pipeline not properly initialized for testing")
            return False

        # Test search functionality
        try:
            search_results = pipeline.search_service.search("vacation policy", top_k=3)
            if search_results and len(search_results) > 0:
                print(f"✅ Search returned {len(search_results)} results for 'vacation policy'")
                return True
            else:
                print("⚠️ Search returned no results - may be expected with test data")
                return True  # This may be expected in some environments
        except Exception as e:
            print(f"❌ Search query failed: {e}")
            return False

    except Exception as e:
        print(f"❌ Query response test failed: {e}")
        return False


def test_health_check():
    """Test RAG pipeline health check"""
    print("🔍 Testing RAG pipeline health check...")

    try:
        # Set environment to enable HF services
        os.environ["OPENROUTER_API_KEY"] = "test"

        from routes.main_routes import get_rag_pipeline

        pipeline = get_rag_pipeline()

        if pipeline is None:
            print("❌ RAG pipeline not initialized for health check")
            return False

        health = pipeline.health_check()
        if health and isinstance(health, dict):
            status = health.get("pipeline", "unknown")
            print(f"✅ Health check completed - Pipeline status: {status}")
            return True
        else:
            print("❌ Health check returned invalid response")
            return False

    except Exception as e:
        print(f"❌ Health check failed: {e}")
        return False


def main():
    """Run all HF deployment validation tests"""
    print("🧪 Testing Complete HF Deployment Functionality")
    print("=" * 60)

    tests = [
        ("File Ingestion", test_file_ingestion),
        ("RAG Pipeline Initialization", test_rag_pipeline_initialization),
        ("Query Response", test_query_response),
        ("Health Check", test_health_check),
    ]

    passed = 0
    for test_name, test_func in tests:
        print(f"\n🔧 Running {test_name} test...")
        try:
            if test_func():
                print(f"✅ {test_name}: PASSED")
                passed += 1
            else:
                print(f"❌ {test_name}: FAILED")
        except Exception as e:
            print(f"💥 {test_name}: ERROR - {e}")
        print()

    print("=" * 60)
    print(f"📊 Results: {passed}/{len(tests)} tests passed")

    if passed == len(tests):
        print("🎉 All tests passed! HF deployment should work correctly.")
        print("📝 The app should now:")
        print("   • Process synthetic_policies files during startup")
        print("   • Create embeddings from policy documents")
        print("   • Initialize RAG pipeline with populated vector store")
        print("   • Return proper responses to policy questions")
        return 0
    else:
        print("⚠️  Some tests failed. Check the issues above.")
        return 1


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