tests/test_document_analyses_integration.py

#!/usr/bin/env python3
"""
Test script for document analyses integration with lawyer prompts
"""

import asyncio
import sys
sys.path.append('.')

from agent_api import CyberLegalAPI, DocumentAnalysis
import logging

logging.basicConfig(level=logging.INFO)

async def test_lawyer_with_documents():
    """Test lawyer agent with document analyses"""
    api = CyberLegalAPI()
    
    # Create sample document analyses
    doc_analyses = [
        DocumentAnalysis(
            file_name="employment_contract.pdf",
            summary="Employment contract with termination clause",
            actors="John Doe (employee), TechCorp (employer)",
            key_details="30-day notice period, 2 months severance, confidentiality clause"
        ),
        DocumentAnalysis(
            file_name="nda_agreement.pdf",
            summary="Non-disclosure agreement",
            actors="TechCorp, ThirdParty Inc",
            key_details="5-year duration, penalties for breach, return of materials"
        )
    ]
    
    # Test request with document analyses
    from agent_api import ChatRequest, Message
    request = ChatRequest(
        message="What are the termination obligations according to the employment contract?",
        conversationHistory=[],
        userType="lawyer",
        jurisdiction="Romania",
        documentAnalyses=doc_analyses
    )
    
    print("Testing lawyer agent with document analyses...")
    print("="*60)
    
    try:
        response = await api.process_request(request)
        print(f"✅ Response generated successfully")
        print(f"Processing time: {response.processing_time:.2f}s")
        print(f"Response: {response.response[:500]}...")
        print("="*60)
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

async def test_lawyer_without_documents():
    """Test lawyer agent without document analyses"""
    api = CyberLegalAPI()
    
    from agent_api import ChatRequest
    request = ChatRequest(
        message="What are the key provisions of GDPR?",
        conversationHistory=[],
        userType="lawyer",
        jurisdiction="Romania",
        documentAnalyses=None
    )
    
    print("\nTesting lawyer agent without document analyses...")
    print("="*60)
    
    try:
        response = await api.process_request(request)
        print(f"✅ Response generated successfully")
        print(f"Processing time: {response.processing_time:.2f}s")
        print(f"Response: {response.response[:500]}...")
        print("="*60)
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

async def test_client_user():
    """Test client user (should ignore document analyses)"""
    api = CyberLegalAPI()
    
    doc_analyses = [
        DocumentAnalysis(
            file_name="sample.pdf",
            summary="Sample document",
            actors="Person A, Person B",
            key_details="Some details"
        )
    ]
    
    from agent_api import ChatRequest
    request = ChatRequest(
        message="What is GDPR?",
        conversationHistory=[],
        userType="client",
        jurisdiction="Romania",
        documentAnalyses=doc_analyses  # Should be ignored for clients
    )
    
    print("\nTesting client user (document analyses should be ignored)...")
    print("="*60)
    
    try:
        response = await api.process_request(request)
        print(f"✅ Response generated successfully")
        print(f"Processing time: {response.processing_time:.2f}s")
        print(f"Response: {response.response[:500]}...")
        print("="*60)
        return True
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

async def main():
    """Run all tests"""
    print("Starting document analyses integration tests...")
    print("="*60)
    
    results = []
    
    # Test 1: Lawyer with documents
    results.append(await test_lawyer_with_documents())
    
    # Test 2: Lawyer without documents
    results.append(await test_lawyer_without_documents())
    
    # Test 3: Client (should ignore documents)
    results.append(await test_client_user())
    
    print("\n" + "="*60)
    print("Test Summary:")
    print(f"Passed: {sum(results)}/{len(results)}")
    print("="*60)
    
    return all(results)

if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)