test_integration.py

#!/usr/bin/env python3
"""
Comprehensive integration test for the Dynamic RAG Bot.
Tests all conversation patterns and flows between main bot and document loader.
"""

from bot import ask_bot
from document_reader import load_document

def parse_fetch_action(action: str) -> tuple[str, str]:
    """Parse fetch action to extract country and question."""
    if not action.startswith("fetch:"):
        return None, None
    
    parts = action[6:].split(":", 2)
    
    if len(parts) >= 2:
        return parts[0], parts[1]
    elif len(parts) == 1:
        return parts[0], "general information"
    
    return None, None

def handle_bot_response(query: str, expected_action=None):
    """Handle a bot response with full integration logic."""
    print(f"User: {query}")
    
    try:
        response = ask_bot(query)
        print(f"[Thinking: {response.thoughts}]")
        print(f"[Action: {response.action}]")
        
        if expected_action:
            actual_type = response.action.split(":")[0] if ":" in response.action else response.action
            if actual_type != expected_action:
                print(f"FAIL: Expected action '{expected_action}', got '{actual_type}'")
                return False
        
        # Handle different actions
        if response.action == "answer":
            print(f"Bot: {response.content}")
            
        elif response.action == "ask_specifics":
            print(f"Bot: {response.content}")
            
        elif response.action.startswith("fetch:"):
            country, question = parse_fetch_action(response.action)
            
            if country and question:
                print(f"[Fetching document for '{country}' about: '{question}']")
                doc_response = load_document(country, question)
                print(f"Bot: {doc_response.format_with_sources()}")
            else:
                print(f"ERROR: Could not parse fetch action: {response.action}")
                
        else:
            print(f"UNKNOWN ACTION: {response.action}")
            if response.content:
                print(f"Bot: {response.content}")
        
        return True
        
    except Exception as e:
        print(f"ERROR: {e}")
        return False

def test_comprehensive_integration():
    """Run comprehensive integration tests."""
    
    print("=" * 80)
    print("COMPREHENSIVE INTEGRATION TEST: Dynamic RAG Bot")
    print("=" * 80)
    
    test_cases = [
        # Basic fact queries - should use 'answer' action
        ("Basic Facts", [
            ("What's the capital of France?", "answer"),
            ("What's the population of Japan?", "answer"), 
            ("What currency does Germany use?", "answer"),
            ("What language do they speak in Italy?", "answer"),
        ]),
        
        # Detailed queries - should use 'fetch' action  
        ("Detailed Information", [
            ("Tell me about French wine regions", "fetch"),
            ("What are the main tourist attractions in Japan?", "fetch"),
            ("Tell me more about Germany", "fetch"),
            ("What is Japanese culture like?", "fetch"),
        ]),
        
        # Edge cases
        ("Edge Cases", [
            ("Tell me about Mexico", "answer"),  # Unknown country
            ("Yes", "ask_specifics"),  # Generic yes response
            ("What are the main industries in Brazil?", "fetch"),  # Economic info
        ]),
        
        # Different question patterns
        ("Question Patterns", [
            ("What is France known for?", "fetch"),  # General cultural question
            ("Tell me about Spanish history", "fetch"),  # Historical info
            ("What's special about Indian cuisine?", "fetch"),  # Cultural detail
        ])
    ]
    
    total_tests = 0
    passed_tests = 0
    
    for category, tests in test_cases:
        print(f"\n{'='*20} {category} {'='*20}")
        
        for query, expected_action in tests:
            print(f"\n{'-'*60}")
            total_tests += 1
            
            success = handle_bot_response(query, expected_action)
            if success:
                passed_tests += 1
                print("[PASS] Test passed")
            else:
                print("[FAIL] Test failed")
    
    # Summary
    print(f"\n{'='*80}")
    print(f"INTEGRATION TEST RESULTS")
    print(f"{'='*80}")
    print(f"Tests passed: {passed_tests}/{total_tests}")
    print(f"Success rate: {(passed_tests/total_tests)*100:.1f}%")
    
    if passed_tests == total_tests:
        print("All tests passed! Integration is working perfectly.")
    else:
        print(f"{total_tests - passed_tests} tests failed. Review integration.")
    
    print(f"{'='*80}")

def test_conversation_flow():
    """Test a realistic conversation flow."""
    
    print(f"\n{'='*80}")
    print("CONVERSATION FLOW TEST")
    print(f"{'='*80}")
    
    conversation = [
        "What's the capital of France?",
        "Tell me about French wine regions", 
        "What's the population of Japan?",
        "What are the main tourist attractions in Japan?"
    ]
    
    print("Testing realistic conversation flow:")
    for i, query in enumerate(conversation, 1):
        print(f"\n--- Turn {i} ---")
        handle_bot_response(query)
    
    print("\n[COMPLETE] Conversation flow test complete")

if __name__ == "__main__":
    test_comprehensive_integration()
    test_conversation_flow()