tests/test_documents_tree.py

#!/usr/bin/env python3
"""
Test the documents_tree integration with the API
"""

import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from structured_outputs.api_models import (
    TreeNode, DocumentsTree, DocumentAnalysis, ChatRequest
)

def test_tree_node_creation():
    """Test creating a tree node"""
    print("🧪 Testing TreeNode creation...")
    
    file_analysis = DocumentAnalysis(
        summary="Test summary",
        actors="Actor 1, Actor 2",
        key_details="Key detail 1, Key detail 2"
    )
    
    file_node = TreeNode(
        name="test.pdf",
        type="file",
        file_path="user-id/test.pdf",
        file_size=100000,
        mime_type="application/pdf",
        created_at="2025-01-15T10:30:00Z",
        analysis=file_analysis
    )
    
    assert file_node.name == "test.pdf"
    assert file_node.type == "file"
    assert file_node.analysis.summary == "Test summary"
    print("✅ TreeNode creation works")

def test_documents_tree_creation():
    """Test creating a documents tree"""
    print("\n🧪 Testing DocumentsTree creation...")
    
    # Create file nodes with analyses
    file1_analysis = DocumentAnalysis(
        summary="Contract summary",
        actors="SCI Martin, SARL Dupont",
        key_details="Durée: 9 ans, Loyer: 3500€/mois"
    )
    
    file1 = TreeNode(
        name="bail-commercial.pdf",
        type="file",
        file_path="user-id/abc123-bail-commercial.pdf",
        file_size=245000,
        mime_type="application/pdf",
        created_at="2025-01-15T10:30:00Z",
        analysis=file1_analysis
    )
    
    file2_analysis = DocumentAnalysis(
        summary="Legal note summary",
        actors="Entreprise XYZ, CNIL",
        key_details="Base légale: intérêt légitime"
    )
    
    file2 = TreeNode(
        name="note-juridique.pdf",
        type="file",
        file_path="user-id/def456-note-juridique.pdf",
        file_size=120000,
        mime_type="application/pdf",
        created_at="2025-02-01T14:00:00Z",
        analysis=file2_analysis
    )
    
    # Create folder with files
    contracts_folder = TreeNode(
        name="Contracts",
        type="folder",
        children=[file1]
    )
    
    # Create root tree
    tree = DocumentsTree(
        name="root",
        type="folder",
        children=[contracts_folder, file2]
    )
    
    assert tree.name == "root"
    assert tree.type == "folder"
    assert len(tree.children) == 2
    assert tree.children[0].name == "Contracts"
    assert tree.children[0].type == "folder"
    assert tree.children[0].children[0].name == "bail-commercial.pdf"
    assert tree.children[1].name == "note-juridique.pdf"
    print("✅ DocumentsTree creation works")

def test_chat_request_with_tree():
    """Test creating a ChatRequest with documents_tree"""
    print("\n🧪 Testing ChatRequest with documents_tree...")
    
    # Create a simple tree
    file_analysis = DocumentAnalysis(
        summary="Test document",
        actors="Actor 1",
        key_details="Detail 1"
    )
    
    file_node = TreeNode(
        name="test.pdf",
        type="file",
        analysis=file_analysis
    )
    
    tree = DocumentsTree(
        children=[file_node]
    )
    
    # Create chat request
    request = ChatRequest(
        clientId="test-client",
        message="What are my documents?",
        userType="lawyer",
        jurisdiction="Romania",
        documents_tree=tree
    )
    
    assert request.clientId == "test-client"
    assert request.message == "What are my documents?"
    assert request.userType == "lawyer"
    assert request.documents_tree is not None
    assert request.documents_tree.children[0].name == "test.pdf"
    print("✅ ChatRequest with documents_tree works")

def test_format_documents_tree():
    """Test the _format_documents_tree method"""
    print("\n🧪 Testing _format_documents_tree...")
    
    # Import the API to access the method
    from agent_api import CyberLegalAPI
    
    # Create a tree
    file_analysis = DocumentAnalysis(
        summary="A very long summary that should be truncated at 100 characters and then show ellipsis",
        actors="Actor 1, Actor 2, Actor 3",
        key_details="Key detail"
    )
    
    file_node = TreeNode(
        name="test.pdf",
        type="file",
        analysis=file_analysis
    )
    
    tree = DocumentsTree(
        children=[file_node]
    )
    
    # Create API instance
    api = CyberLegalAPI()
    
    # Format the tree
    formatted = api._format_documents_tree(tree)
    
    print("📄 Formatted tree:")
    print(formatted)
    
    assert "test.pdf" in formatted
    assert "summary:" in formatted
    assert "actors:" in formatted
    assert "key_details:" in formatted
    print("✅ _format_documents_tree works")

def test_extract_flat_documents():
    """Test the _extract_flat_documents method"""
    print("\n🧪 Testing _extract_flat_documents...")
    
    # Import the API to access the method
    from agent_api import CyberLegalAPI
    
    # Create a tree with multiple files in folders
    file1_analysis = DocumentAnalysis(
        summary="File 1 summary",
        actors="Actor 1",
        key_details="Detail 1"
    )
    
    file1 = TreeNode(
        name="file1.pdf",
        type="file",
        analysis=file1_analysis
    )
    
    file2_analysis = DocumentAnalysis(
        summary="File 2 summary",
        actors="Actor 2",
        key_details="Detail 2"
    )
    
    file2 = TreeNode(
        name="file2.pdf",
        type="file",
        analysis=file2_analysis
    )
    
    folder = TreeNode(
        name="Folder",
        type="folder",
        children=[file1]
    )
    
    tree = DocumentsTree(
        children=[folder, file2]
    )
    
    # Create API instance
    api = CyberLegalAPI()
    
    # Extract flat documents
    flat_docs = api._extract_flat_documents(tree)
    
    print(f"📄 Extracted {len(flat_docs)} documents:")
    for doc in flat_docs:
        print(f"  - {doc['file_name']}: {doc['summary']}")
    
    assert len(flat_docs) == 2
    assert flat_docs[0]['file_name'] == 'file1.pdf'
    assert flat_docs[0]['summary'] == 'File 1 summary'
    assert flat_docs[1]['file_name'] == 'file2.pdf'
    assert flat_docs[1]['summary'] == 'File 2 summary'
    print("✅ _extract_flat_documents works")

if __name__ == "__main__":
    print("🚀 Running documents_tree integration tests\n")
    print("=" * 80)
    
    try:
        test_tree_node_creation()
        test_documents_tree_creation()
        test_chat_request_with_tree()
        test_format_documents_tree()
        test_extract_flat_documents()
        
        print("\n" + "=" * 80)
        print("✅ All tests passed!")
        print("=" * 80)
        
    except Exception as e:
        print("\n" + "=" * 80)
        print(f"❌ Test failed: {e}")
        print("=" * 80)
        import traceback
        traceback.print_exc()
        sys.exit(1)