test_setup.py

#!/usr/bin/env python3
"""
Lineage Graph Extractor - Setup Test Script

This script tests your local setup to ensure everything is configured correctly.

Usage:
    python test_setup.py
"""

import os
import sys
from pathlib import Path

def test_python_version():
    """Test Python version"""
    print("Testing Python version...")
    version = sys.version_info
    if version.major >= 3 and version.minor >= 9:
        print(f"✓ Python {version.major}.{version.minor}.{version.micro} (OK)")
        return True
    else:
        print(f"✗ Python {version.major}.{version.minor}.{version.micro} (Need 3.9+)")
        return False

def test_dependencies():
    """Test if required dependencies are installed"""
    print("\nTesting dependencies...")
    
    dependencies = {
        "anthropic": "Anthropic API client",
        "dotenv": "Environment variable loader (python-dotenv)"
    }
    
    all_installed = True
    for module, description in dependencies.items():
        try:
            if module == "dotenv":
                __import__("dotenv")
            else:
                __import__(module)
            print(f"✓ {description}")
        except ImportError:
            print(f"✗ {description} (not installed)")
            all_installed = False
    
    if not all_installed:
        print("\nInstall missing dependencies with:")
        print("  pip install -r requirements.txt")
    
    return all_installed

def test_env_file():
    """Test if .env file exists and has required variables"""
    print("\nTesting environment configuration...")
    
    if not Path(".env").exists():
        print("✗ .env file not found")
        print("  Copy .env.example to .env and add your API keys")
        return False
    
    print("✓ .env file exists")
    
    # Try to load it
    try:
        from dotenv import load_dotenv
        load_dotenv()
        
        api_key = os.getenv("ANTHROPIC_API_KEY")
        if not api_key or api_key == "your_anthropic_api_key_here":
            print("✗ ANTHROPIC_API_KEY not set or still has default value")
            print("  Edit .env and add your actual Anthropic API key")
            return False
        
        print("✓ ANTHROPIC_API_KEY is set")
        return True
    except Exception as e:
        print(f"✗ Error loading .env: {e}")
        return False

def test_agent_files():
    """Test if agent configuration files exist"""
    print("\nTesting agent configuration files...")
    
    required_files = [
        "memories/agent.md",
        "memories/tools.json",
        "memories/subagents/metadata_parser/agent.md",
        "memories/subagents/metadata_parser/tools.json",
        "memories/subagents/graph_visualizer/agent.md",
        "memories/subagents/graph_visualizer/tools.json"
    ]
    
    all_exist = True
    for file_path in required_files:
        if Path(file_path).exists():
            print(f"✓ {file_path}")
        else:
            print(f"✗ {file_path} (missing)")
            all_exist = False
    
    return all_exist

def test_api_connection():
    """Test connection to Anthropic API"""
    print("\nTesting Anthropic API connection...")
    
    try:
        from anthropic import Anthropic
        from dotenv import load_dotenv
        
        load_dotenv()
        
        api_key = os.getenv("ANTHROPIC_API_KEY")
        if not api_key:
            print("✗ API key not found")
            return False
        
        client = Anthropic(api_key=api_key)
        
        # Make a simple test request
        response = client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=100,
            messages=[{
                "role": "user",
                "content": "Hello"
            }]
        )
        
        print("✓ API connection successful")
        print(f"  Model: {response.model}")
        print(f"  Response: {response.content[0].text[:50]}...")
        return True
        
    except Exception as e:
        print(f"✗ API connection failed: {e}")
        return False

def test_agent_functionality():
    """Test basic agent functionality"""
    print("\nTesting agent functionality...")
    
    try:
        from anthropic import Anthropic
        from dotenv import load_dotenv
        
        load_dotenv()
        
        client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
        
        # Load agent configuration
        with open("memories/agent.md", "r") as f:
            system_prompt = f.read()
        
        print("✓ Agent configuration loaded")
        
        # Test agent response
        response = client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=500,
            system=system_prompt,
            messages=[{
                "role": "user",
                "content": "What types of metadata sources can you extract lineage from?"
            }]
        )
        
        print("✓ Agent responds correctly")
        print(f"  Response preview: {response.content[0].text[:100]}...")
        return True
        
    except Exception as e:
        print(f"✗ Agent test failed: {e}")
        return False

def main():
    """Run all tests"""
    print("=" * 60)
    print("Lineage Graph Extractor - Setup Test")
    print("=" * 60)
    
    results = {
        "Python version": test_python_version(),
        "Dependencies": test_dependencies(),
        "Environment file": test_env_file(),
        "Agent files": test_agent_files(),
        "API connection": test_api_connection(),
        "Agent functionality": test_agent_functionality()
    }
    
    print("\n" + "=" * 60)
    print("Test Summary")
    print("=" * 60)
    
    for test_name, passed in results.items():
        status = "✓ PASS" if passed else "✗ FAIL"
        print(f"{test_name:.<40} {status}")
    
    all_passed = all(results.values())
    
    print("\n" + "=" * 60)
    if all_passed:
        print("✓ All tests passed! Your setup is ready.")
        print("\nNext steps:")
        print("  1. Try the integration example: python integration_example.py")
        print("  2. Read the README.md for usage examples")
        print("  3. Extract your first lineage!")
    else:
        print("✗ Some tests failed. Please fix the issues above.")
        print("\nCommon fixes:")
        print("  - Install dependencies: pip install -r requirements.txt")
        print("  - Copy .env.example to .env and add your API key")
        print("  - Verify all files are present")
    print("=" * 60)
    
    return 0 if all_passed else 1

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