test_file_download.py

#!/usr/bin/env python3
"""
Test file download functionality
"""

import os
from dotenv import load_dotenv
from app import BasicAgent

load_dotenv()

def test_file_download():
    """Test questions with file URLs"""
    
    agent = BasicAgent()
    api_key = os.getenv("ANTHROPIC_API_KEY")
    if not api_key:
        print("Error: ANTHROPIC_API_KEY not found")
        return
    
    agent.set_api_key(api_key)
    
    # Test cases with file URLs (these are hypothetical)
    test_cases = [
        {
            "question": "What is the total sales from the Excel file at https://example.com/sales.xlsx?",
            "type": "excel_url"
        },
        {
            "question": "How many times does 'therefore' appear in the PDF at https://example.com/document.pdf?",
            "type": "pdf_url"
        },
        {
            "question": "The attached Excel file contains sales data. What is the total?",
            "type": "no_url"
        }
    ]
    
    for i, test in enumerate(test_cases, 1):
        print(f"\nTest {i} ({test['type']}):")
        print(f"Question: {test['question']}")
        
        try:
            answer = agent(test['question'])
            print(f"Answer: {answer}")
            
            if test['type'] == 'no_url' and "unable to determine" in answer.lower():
                print("✅ Correctly identified missing file")
            elif test['type'] in ['excel_url', 'pdf_url']:
                if "failed to download" in answer.lower():
                    print("⚠️ URL not accessible (expected for example.com)")
                else:
                    print("✅ Attempted to process URL")
                    
        except Exception as e:
            print(f"Error: {e}")

if __name__ == "__main__":
    test_file_download()