unit_tests/test_integration.py

#!/usr/bin/env python3
"""Integration test for the custom parser with ReActAgent."""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from llama_index.core.agent import ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.core.llms.mock import MockLLM
from llama_index.core.memory import ChatMemoryBuffer
from app.core.custom_parser import VerbatimOutputParser, VERBATIM_RESPONSE_CONFIG


def mock_booking_tool(title: str, time: str) -> str:
    """Mock booking tool that returns a realistic booking confirmation."""
    return """Perfect! ✅ **Meeting confirmed: licensing**
📅 **Tuesday, January 14, 2025 at 2:30 PM** (15 minutes)
<div style="background: #e8f5e9; padding: 10px; border-radius: 6px; margin: 8px 0;"><strong>🎥 Google Meet:</strong> <a href="https://meet.google.com/abc-def-ghi" style="color: #1976d2; font-weight: bold;">Join here</a></div>
<div style="background: #f5f5f5; padding: 10px; border-radius: 6px; margin: 8px 0; font-family: monospace;"><strong>📋 Meeting IDs:</strong><br>Meeting ID: 0114-1430-15m<br>Calendar ID: abc123def456</div>
<div style="color: #2196f3; margin: 8px 0;"><strong>📧 Email invitations are being sent in the background...</strong></div>
<div id="booking-success" style="display:none;">BOOKING_COMPLETE</div>"""


def test_agent_with_custom_parser():
    """Test that the ReActAgent uses the custom parser correctly."""
    print("🔧 Testing ReActAgent with custom parser...")

    # Create mock LLM that simulates booking flow
    mock_responses = [
        "I need to book this meeting. Let me call the booking tool.\n\nAction: book_meeting\nAction Input: {\"title\": \"licensing\", \"time\": \"2:30 PM\"}"
    ]

    mock_llm = MockLLM(max_tokens=None)
    mock_llm._responses = mock_responses

    # Create tool
    booking_tool = FunctionTool.from_defaults(
        fn=mock_booking_tool,
        name="book_meeting",
        description="Book a meeting"
    )

    # Create custom parser
    custom_parser = VerbatimOutputParser(verbatim_patterns=VERBATIM_RESPONSE_CONFIG['patterns'])

    # Create agent with custom parser
    memory = ChatMemoryBuffer(token_limit=3000)
    agent = ReActAgent(
        tools=[booking_tool],
        llm=mock_llm,
        memory=memory,
        verbose=True,
        output_parser=custom_parser
    )

    print("✅ Agent created successfully with custom parser")
    print(f"   Agent type: {type(agent)}")
    print(f"   Custom parser provided to agent initialization")

    # Test that the parser can be called directly
    test_output = """Observation: Perfect! ✅ **Meeting confirmed: licensing**
📅 **Tuesday, January 14, 2025 at 2:30 PM** (15 minutes)"""

    result = custom_parser.parse(test_output)
    print(f"   Direct parser test: {type(result).__name__}")

    return True


def test_parser_configuration():
    """Test that the parser configuration is properly loaded."""
    print("⚙️  Testing parser configuration...")

    config = VERBATIM_RESPONSE_CONFIG

    print(f"   Enabled: {config['enabled']}")
    print(f"   Tools count: {len(config['tools'])}")
    print(f"   Patterns count: {len(config['patterns'])}")
    print(f"   Fallback on error: {config['fallback_on_error']}")

    # Test that all expected patterns are present
    expected_patterns = [
        'booking_confirmation',
        'cancellation_success',
        'google_meet',
        'authentication_error'
    ]

    missing_patterns = [p for p in expected_patterns if p not in config['patterns']]
    if missing_patterns:
        print(f"❌ Missing patterns: {missing_patterns}")
        return False

    print("✅ All expected patterns present")
    return True


def main():
    """Run integration tests."""
    print("🧪 Integration Tests for Custom Output Parser\n")

    print("=" * 60)
    success1 = test_parser_configuration()

    print("\n" + "=" * 60)
    success2 = test_agent_with_custom_parser()

    print("\n" + "=" * 60)
    if success1 and success2:
        print("🎉 All integration tests passed!")
        return 0
    else:
        print("❌ Some tests failed!")
        return 1


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