endpoint_test_script.py

#!/usr/bin/env python3
"""
Comprehensive test script for Flask RAG API endpoints
Run this script to test all your endpoints systematically
"""

import requests
import json
import time
import os
from pathlib import Path

# Configuration
BASE_URL = "http://localhost:5001"
TEST_FILE_PATH = "test_document.txt"

class APITester:
    def __init__(self, base_url=BASE_URL):
        self.base_url = base_url
        self.session = requests.Session()
        self.conversation_id = None
        self.test_results = []
        
    def log_result(self, test_name, success, response_data=None, error=None):
        """Log test results"""
        result = {
            "test": test_name,
            "success": success,
            "timestamp": time.strftime("%H:%M:%S"),
            "response": response_data,
            "error": str(error) if error else None
        }
        self.test_results.append(result)
        
        status = "✅ PASS" if success else "❌ FAIL"
        print(f"{status} {test_name}")
        if error:
            print(f"    Error: {error}")
        if response_data and isinstance(response_data, dict):
            print(f"    Response: {json.dumps(response_data, indent=2)[:200]}...")
        print()
    
    def create_test_file(self):
        """Create a test document for upload"""
        test_content = """
        # Test Document for RAG System
        
        This is a test document for the RAG (Retrieval-Augmented Generation) system.
        
        ## Key Information
        - This document contains test data
        - The RAG system should be able to search and retrieve this content
        - OpenAI embeddings are used for vector similarity search
        - Pinecone is used as the vector database
        
        ## Technical Details
        The system uses:
        1. OpenAI text-embedding-3-small for embeddings (1536 dimensions)
        2. Pinecone serverless for vector storage
        3. Flask for the API backend
        4. Chunking strategy with 500 character chunks and 50 character overlap
        
        ## Sample Questions
        You can ask questions like:
        - "What embedding model is used?"
        - "How many dimensions do the embeddings have?"
        - "What vector database is used?"
        """
        
        with open(TEST_FILE_PATH, 'w') as f:
            f.write(test_content)
        print(f"📝 Created test file: {TEST_FILE_PATH}")
    
    def test_health_check(self):
        """Test the health check endpoint"""
        try:
            response = self.session.get(f"{self.base_url}/health")
            if response.status_code == 200:
                data = response.json()
                self.log_result("Health Check", True, data)
                return True
            else:
                self.log_result("Health Check", False, error=f"Status: {response.status_code}")
                return False
        except Exception as e:
            self.log_result("Health Check", False, error=e)
            return False
    
    def test_services_info(self):
        """Test the services info endpoint"""
        try:
            response = self.session.get(f"{self.base_url}/services/info")
            if response.status_code == 200:
                data = response.json()
                self.log_result("Services Info", True, data)
                return True
            else:
                self.log_result("Services Info", False, error=f"Status: {response.status_code}")
                return False
        except Exception as e:
            self.log_result("Services Info", False, error=e)
            return False
    
    def test_upload_document(self):
        """Test document upload to RAG system"""
        try:
            if not os.path.exists(TEST_FILE_PATH):
                self.create_test_file()
            
            with open(TEST_FILE_PATH, 'rb') as f:
                files = {'file': (TEST_FILE_PATH, f, 'text/plain')}
                data = {'metadata': json.dumps({"source": "test", "type": "documentation"})}
                
                response = self.session.post(
                    f"{self.base_url}/rag/upload",
                    files=files,
                    data=data
                )
            
            if response.status_code == 200:
                data = response.json()
                self.log_result("Document Upload", True, data)
                return data.get('document_id')
            else:
                self.log_result("Document Upload", False, error=f"Status: {response.status_code}, Response: {response.text}")
                return None
        except Exception as e:
            self.log_result("Document Upload", False, error=e)
            return None
    
    def test_list_documents(self):
        """Test listing uploaded documents"""
        try:
            response = self.session.get(f"{self.base_url}/rag/documents")
            if response.status_code == 200:
                data = response.json()
                self.log_result("List Documents", True, data)
                return True
            else:
                self.log_result("List Documents", False, error=f"Status: {response.status_code}")
                return False
        except Exception as e:
            self.log_result("List Documents", False, error=e)
            return False
    
    def test_rag_search(self):
        """Test RAG search functionality"""
        try:
            query_data = {
                "query": "What embedding model is used in the system?",
                "top_k": 3
            }
            
            response = self.session.post(
                f"{self.base_url}/rag/search",
                json=query_data,
                headers={'Content-Type': 'application/json'}
            )
            
            if response.status_code == 200:
                data = response.json()
                self.log_result("RAG Search", True, data)
                return True
            else:
                self.log_result("RAG Search", False, error=f"Status: {response.status_code}, Response: {response.text}")
                return False
        except Exception as e:
            self.log_result("RAG Search", False, error=e)
            return False
    
    def test_chat_with_rag(self):
        """Test chat endpoint with RAG context"""
        try:
            chat_data = {
                "message": "What vector database is used in this system?",
                "use_rag": True,
                "conversation_id": self.conversation_id
            }
            
            response = self.session.post(
                f"{self.base_url}/chat",
                json=chat_data,
                headers={'Content-Type': 'application/json'}
            )
            
            if response.status_code == 200:
                data = response.json()
                self.conversation_id = data.get('conversation_id')  # Store for future tests
                self.log_result("Chat with RAG", True, data)
                return True
            else:
                self.log_result("Chat with RAG", False, error=f"Status: {response.status_code}, Response: {response.text}")
                return False
        except Exception as e:
            self.log_result("Chat with RAG", False, error=e)
            return False
    
    def test_chat_without_rag(self):
        """Test chat endpoint without RAG context"""
        try:
            chat_data = {
                "message": "Hello, how are you?",
                "use_rag": False,
                "conversation_id": self.conversation_id
            }
            
            response = self.session.post(
                f"{self.base_url}/chat",
                json=chat_data,
                headers={'Content-Type': 'application/json'}
            )
            
            if response.status_code == 200:
                data = response.json()
                self.log_result("Chat without RAG", True, data)
                return True
            else:
                self.log_result("Chat without RAG", False, error=f"Status: {response.status_code}, Response: {response.text}")
                return False
        except Exception as e:
            self.log_result("Chat without RAG", False, error=e)
            return False
    
    def test_conversation_history(self):
        """Test getting conversation history"""
        if not self.conversation_id:
            self.log_result("Conversation History", False, error="No conversation ID available")
            return False
        
        try:
            response = self.session.get(f"{self.base_url}/conversation/{self.conversation_id}")
            
            if response.status_code == 200:
                data = response.json()
                self.log_result("Conversation History", True, data)
                return True
            else:
                self.log_result("Conversation History", False, error=f"Status: {response.status_code}")
                return False
        except Exception as e:
            self.log_result("Conversation History", False, error=e)
            return False
    
    def test_tts_synthesis(self):
        """Test text-to-speech synthesis"""
        try:
            tts_data = {
                "text": "Hello, this is a test of the text to speech system.",
                "voice": "alloy"
            }
            
            response = self.session.post(
                f"{self.base_url}/tts/synthesize",
                json=tts_data,
                headers={'Content-Type': 'application/json'}
            )
            
            if response.status_code == 200:
                # TTS returns audio data, so we just check the response headers
                content_type = response.headers.get('content-type', '')
                success = 'audio' in content_type or len(response.content) > 1000
                self.log_result("TTS Synthesis", success, {"content_length": len(response.content), "content_type": content_type})
                return success
            else:
                self.log_result("TTS Synthesis", False, error=f"Status: {response.status_code}")
                return False
        except Exception as e:
            self.log_result("TTS Synthesis", False, error=e)
            return False
    
    def run_all_tests(self):
        """Run all tests in sequence"""
        print("🚀 Starting comprehensive API endpoint tests")
        print("=" * 60)
        
        # Basic connectivity tests
        print("📡 Testing basic connectivity...")
        health_ok = self.test_health_check()
        info_ok = self.test_services_info()
        
        if not health_ok:
            print("❌ Health check failed - stopping tests")
            return self.print_summary()
        
        # Document and RAG tests
        print("📚 Testing document and RAG functionality...")
        doc_id = self.test_upload_document()
        self.test_list_documents()
        
        if doc_id:
            # Wait a moment for vector indexing
            print("⏳ Waiting for document indexing...")
            time.sleep(2)
            
            self.test_rag_search()
            self.test_chat_with_rag()
        else:
            print("⚠️  Skipping RAG tests due to upload failure")
        
        # Chat tests
        print("💬 Testing chat functionality...")
        self.test_chat_without_rag()
        self.test_conversation_history()
        
        # Audio tests
        print("🔊 Testing audio functionality...")
        self.test_tts_synthesis()
        
        # Note: STT testing would require audio file upload
        print("ℹ️  STT testing requires audio file - skipping for now")
        
        return self.print_summary()
    
    def print_summary(self):
        """Print test summary"""
        print("\n" + "=" * 60)
        print("📊 TEST SUMMARY")
        print("=" * 60)
        
        total_tests = len(self.test_results)
        passed_tests = sum(1 for result in self.test_results if result['success'])
        failed_tests = total_tests - passed_tests
        
        print(f"Total Tests: {total_tests}")
        print(f"✅ Passed: {passed_tests}")
        print(f"❌ Failed: {failed_tests}")
        print(f"Success Rate: {(passed_tests/total_tests)*100:.1f}%")
        
        if failed_tests > 0:
            print("\n❌ Failed Tests:")
            for result in self.test_results:
                if not result['success']:
                    print(f"   - {result['test']}: {result['error']}")
        
        print("\n🧹 Cleaning up...")
        if os.path.exists(TEST_FILE_PATH):
            os.remove(TEST_FILE_PATH)
            print(f"   Removed test file: {TEST_FILE_PATH}")
        
        return passed_tests == total_tests

def main():
    """Main test function"""
    print("🧪 Flask RAG API Endpoint Tester")
    print(f"🎯 Target: {BASE_URL}")
    print("=" * 60)
    
    # Check if server is running
    try:
        response = requests.get(f"{BASE_URL}/health", timeout=5)
        print("✅ Server is responding")
    except requests.exceptions.RequestException as e:
        print(f"❌ Server is not responding: {e}")
        print("💡 Make sure your Flask app is running on http://localhost:5001")
        return False
    
    # Run tests
    tester = APITester()
    success = tester.run_all_tests()
    
    if success:
        print("🎉 All tests passed!")
    else:
        print("⚠️  Some tests failed - check the logs above")
    
    return success

if __name__ == "__main__":
    main()