test_integration.py

"""Comprehensive integration test for RAG Service"""
import sys
import asyncio
from pathlib import Path


def test_configuration():
    """Test configuration and API keys"""
    print("=" * 60)
    print("1. Testing Configuration")
    print("=" * 60)
    
    try:
        from app.core.config import get_settings
        settings = get_settings()
        
        print(f"✓ Settings loaded")
        print(f"  - Database URL: {settings.database_url}")
        print(f"  - ChromaDB Path: {settings.chroma_db_path}")
        print(f"  - Embedding Model: {settings.gemini_embedding_model}")
        print(f"  - Embedding Dimension: {settings.embedding_dimension}")
        print(f"  - Port: {settings.port}")
        print(f"  - File Service URL: {settings.file_service_url}")
        
        # Check API key
        if settings.gemini_api_key:
            key_preview = settings.gemini_api_key[:10] + "..." if len(settings.gemini_api_key) > 10 else "***"
            print(f"✓ Gemini API Key configured: {key_preview}")
        else:
            print("⚠ WARNING: Gemini API Key NOT configured!")
            return False
        
        print("\n✅ Configuration tests passed!\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Configuration test failed: {e}\n")
        return False


def test_database_initialization():
    """Test database initialization"""
    print("=" * 60)
    print("2. Testing Database Initialization")
    print("=" * 60)
    
    try:
        from app.db.database import init_db, init_chroma, get_chroma_collection
        
        # Initialize databases
        print("Initializing SQLite...")
        init_db()
        print("✓ SQLite initialized")
        
        print("Initializing ChromaDB...")
        init_chroma()
        print("✓ ChromaDB initialized")
        
        # Check ChromaDB
        collection = get_chroma_collection()
        count = collection.count()
        print(f"✓ ChromaDB collection count: {count}")
        
        print("\n✅ Database initialization tests passed!\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Database initialization test failed: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_models():
    """Test SQLAlchemy models"""
    print("=" * 60)
    print("3. Testing Models")
    print("=" * 60)
    
    try:
        from app.models.document import Document
        from app.db.database import SessionLocal
        
        # Create session
        db = SessionLocal()
        
        # Query (should be empty initially)
        count = db.query(Document).count()
        print(f"✓ Document model works")
        print(f"  - Current document count: {count}")
        
        db.close()
        
        print("\n✅ Model tests passed!\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Model test failed: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_schemas():
    """Test Pydantic schemas"""
    print("=" * 60)
    print("4. Testing Schemas")
    print("=" * 60)
    
    try:
        from app.schemas.document import (
            DocumentCreate, DocumentSchema, DocumentUploadResponse,
            DocumentDeleteRequest, DocumentDeleteResponse
        )
        from app.schemas.retrieval import (
            RetrievalRequest, RetrievalResponse, RetrievalResult
        )
        from datetime import datetime
        
        # Test document schemas
        doc_create = DocumentCreate(
            user_id=1,
            file_name="test.pdf",
            file_path="/test/path.pdf",
            chunk_count=5
        )
        print("✓ DocumentCreate schema works")
        
        # Test retrieval schemas
        retrieval_req = RetrievalRequest(
            query="test query",
            top_k=5
        )
        print("✓ RetrievalRequest schema works")
        
        print("\n✅ Schema tests passed!\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Schema test failed: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_embedding_service():
    """Test embedding service with real API"""
    print("=" * 60)
    print("5. Testing Embedding Service (Real API)")
    print("=" * 60)
    
    try:
        from app.services.embedding_service import EmbeddingService
        import numpy as np
        
        service = EmbeddingService()
        print(f"✓ EmbeddingService initialized")
        print(f"  - Model: {service.model}")
        print(f"  - Dimension: {service.dimension}")
        
        # Test query embedding
        print("\nTesting query embedding...")
        query = "What is machine learning?"
        query_embedding = service.embed_query(query)
        
        print(f"✓ Query embedding generated")
        print(f"  - Shape: {query_embedding.shape}")
        print(f"  - Type: {type(query_embedding)}")
        print(f"  - Norm: {np.linalg.norm(query_embedding):.6f} (should be ~1.0)")
        
        # Test document embeddings
        print("\nTesting document embeddings...")
        docs = [
            "Machine learning is a subset of AI",
            "Python is a programming language"
        ]
        doc_embeddings = service.embed_documents(docs)
        
        print(f"✓ Document embeddings generated")
        print(f"  - Shape: {doc_embeddings.shape}")
        print(f"  - Count: {len(doc_embeddings)} embeddings")
        print(f"  - First norm: {np.linalg.norm(doc_embeddings[0]):.6f}")
        print(f"  - Second norm: {np.linalg.norm(doc_embeddings[1]):.6f}")
        
        print("\n✅ Embedding service tests passed!\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Embedding service test failed: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_chunking_service():
    """Test chunking service"""
    print("=" * 60)
    print("6. Testing Chunking Service")
    print("=" * 60)
    
    try:
        from app.services.chunking_service import ChunkingService
        
        service = ChunkingService()
        print(f"✓ ChunkingService initialized")
        print(f"  - Chunk size: {service.chunk_size}")
        print(f"  - Overlap: {service.chunk_overlap}")
        
        # Test chunking
        text = """
        Machine learning is a method of data analysis that automates analytical model building.
        It is a branch of artificial intelligence based on the idea that systems can learn from data,
        identify patterns and make decisions with minimal human intervention.
        
        Deep learning is a subset of machine learning that uses neural networks with multiple layers.
        These neural networks attempt to simulate the behavior of the human brain allowing it to
        learn from large amounts of data.
        """ * 5  # Make it longer to test chunking
        
        chunks = service.chunk_text(text)
        print(f"✓ Text chunked")
        print(f"  - Input length: {len(text)} chars")
        print(f"  - Chunks created: {len(chunks)}")
        print(f"  - First chunk length: {len(chunks[0])} chars")
        print(f"  - Last chunk length: {len(chunks[-1])} chars")
        
        print("\n✅ Chunking service tests passed!\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Chunking service test failed: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_retrieval_service():
    """Test retrieval service"""
    print("=" * 60)
    print("7. Testing Retrieval Service")
    print("=" * 60)
    
    try:
        from app.services.retrieval_service import RetrievalService
        from app.db.database import get_chroma_collection
        
        service = RetrievalService()
        print(f"✓ RetrievalService initialized")
        
        # Get collection
        collection = get_chroma_collection()
        count = collection.count()
        
        if count == 0:
            print(f"⚠ No documents in ChromaDB yet (count: {count})")
            print("  - Retrieval will work once documents are uploaded")
        else:
            print(f"✓ ChromaDB has {count} chunks")
            
            # Test retrieval
            query = "machine learning"
            results = service.retrieve_relevant_chunks(query, top_k=3, chroma_collection=collection)
            
            print(f"✓ Retrieval successful")
            print(f"  - Query: '{results.query}'")
            print(f"  - Results found: {results.total_results}")
            
            for i, result in enumerate(results.results[:2]):
                print(f"  - Result {i+1}: score={result.score:.4f}, text={result.text[:50]}...")
        
        print("\n✅ Retrieval service tests passed!\n")
        return True
        
    except Exception as e:
        print(f"\n❌ Retrieval service test failed: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def test_api_routes():
    """Test API routes"""
    print("=" * 60)
    print("8. Testing API Routes")
    print("=" * 60)
    
    try:
        from app.main import app
        
        # Check routes
        routes = [route.path for route in app.routes]
        
        expected = ["/", "/health", "/rag/documents", "/rag/retrieval"]
        
        for route in expected:
            found = any(route in r for r in routes)
            if found:
                print(f"✓ Route exists: {route}")
            else:
                print(f"✗ Route missing: {route}")
                return False
        
        print("\n✅ API routes tests passed!\n")
        return True
        
    except Exception as e:
        print(f"\n❌ API routes test failed: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def main():
    """Run all tests"""
    print("\n")
    print("=" * 60)
    print("RAG SERVICE - COMPREHENSIVE INTEGRATION TEST")
    print("=" * 60)
    print("\n")
    
    results = []
    
    # Run tests
    results.append(("Configuration", test_configuration()))
    results.append(("Database Initialization", test_database_initialization()))
    results.append(("Models", test_models()))
    results.append(("Schemas", test_schemas()))
    results.append(("Embedding Service (API)", test_embedding_service()))
    results.append(("Chunking Service", test_chunking_service()))
    results.append(("Retrieval Service", test_retrieval_service()))
    results.append(("API Routes", test_api_routes()))
    
    # Summary
    print("\n")
    print("=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)
    
    for name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{status} - {name}")
    
    all_passed = all(result for _, result in results)
    
    print("\n")
    print("=" * 60)
    
    if all_passed:
        print("🎉 ALL TESTS PASSED!")
        print("=" * 60)
        print("\nRAG Service is ready for production!")
        print("\nNext steps:")
        print("1. Start the service: python -m uvicorn app.main:app --reload --port 3005")
        print("2. View API docs: http://localhost:3005/docs")
        print("3. Test endpoints with Postman or curl")
        print("\nEndpoints:")
        print("  [ADMIN] POST   /rag/documents      - Upload documents")
        print("  [ADMIN] GET    /rag/documents      - List documents")
        print("  [ADMIN] DELETE /rag/documents      - Delete documents")
        print("  [ALL]   GET    /rag/retrieval      - Retrieve chunks")
        return 0
    else:
        print("⚠ SOME TESTS FAILED")
        print("=" * 60)
        print("\nPlease check the errors above and fix them.")
        return 1


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