claude-scientific-skills/scientific-skills/scientific-schematics/scripts/test_ai_generation.py

#!/usr/bin/env python3
"""
Test script to verify AI generation implementation.

This script performs dry-run tests without making actual API calls.
It verifies:
1. Script structure and imports
2. Class initialization
3. Method signatures
4. Error handling
5. Command-line interface

Usage:
    python test_ai_generation.py
"""

import sys
import os
from pathlib import Path

# Add scripts directory to path
scripts_dir = Path(__file__).parent / "scripts"
sys.path.insert(0, str(scripts_dir))

def test_imports():
    """Test that all required modules can be imported."""
    print("Testing imports...")
    try:
        from generate_schematic_ai import ScientificSchematicGenerator
        print("✓ generate_schematic_ai imports successfully")
        return True
    except ImportError as e:
        print(f"✗ Import failed: {e}")
        return False

def test_class_structure():
    """Test class initialization and structure."""
    print("\nTesting class structure...")
    try:
        from generate_schematic_ai import ScientificSchematicGenerator
        
        # Test initialization with dummy key
        generator = ScientificSchematicGenerator(api_key="test_key", verbose=False)
        print("✓ Class initializes successfully")
        
        # Check required methods exist
        required_methods = [
            'generate_image',
            'review_image',
            'improve_prompt',
            'generate_iterative'
        ]
        
        for method in required_methods:
            if not hasattr(generator, method):
                print(f"✗ Missing method: {method}")
                return False
            print(f"✓ Method exists: {method}")
        
        # Check attributes
        if not hasattr(generator, 'api_key'):
            print("✗ Missing attribute: api_key")
            return False
        print("✓ Attribute exists: api_key")
        
        if not hasattr(generator, 'image_model'):
            print("✗ Missing attribute: image_model")
            return False
        print(f"✓ Image model: {generator.image_model}")
        
        if not hasattr(generator, 'review_model'):
            print("✗ Missing attribute: review_model")
            return False
        print(f"✓ Review model: {generator.review_model}")
        
        return True
    except Exception as e:
        print(f"✗ Class structure test failed: {e}")
        return False

def test_error_handling():
    """Test error handling for missing API key."""
    print("\nTesting error handling...")
    try:
        from generate_schematic_ai import ScientificSchematicGenerator
        
        # Clear environment variable
        old_key = os.environ.get("OPENROUTER_API_KEY")
        if old_key:
            del os.environ["OPENROUTER_API_KEY"]
        
        # Try to initialize without key
        try:
            generator = ScientificSchematicGenerator()
            print("✗ Should have raised ValueError for missing API key")
            return False
        except ValueError as e:
            if "OPENROUTER_API_KEY" in str(e):
                print("✓ Correctly raises ValueError for missing API key")
            else:
                print(f"✗ Wrong error message: {e}")
                return False
        
        # Restore environment variable
        if old_key:
            os.environ["OPENROUTER_API_KEY"] = old_key
        
        return True
    except Exception as e:
        print(f"✗ Error handling test failed: {e}")
        return False

def test_wrapper_script():
    """Test wrapper script structure."""
    print("\nTesting wrapper script...")
    try:
        import generate_schematic
        print("✓ generate_schematic imports successfully")
        
        # Check main functions exist
        if not hasattr(generate_schematic, 'main'):
            print("✗ Missing function: main")
            return False
        print("✓ Function exists: main")
        
        return True
    except Exception as e:
        print(f"✗ Wrapper script test failed: {e}")
        return False

def test_prompt_engineering():
    """Test prompt construction."""
    print("\nTesting prompt engineering...")
    try:
        from generate_schematic_ai import ScientificSchematicGenerator
        
        generator = ScientificSchematicGenerator(api_key="test_key", verbose=False)
        
        # Test improve_prompt method
        original = "Create a flowchart"
        critique = "Add more spacing between boxes"
        improved = generator.improve_prompt(original, critique, 2)
        
        if not improved:
            print("✗ improve_prompt returned empty string")
            return False
        
        if original not in improved:
            print("✗ Improved prompt doesn't include original")
            return False
        
        if critique not in improved:
            print("✗ Improved prompt doesn't include critique")
            return False
        
        if "ITERATION 2" not in improved:
            print("✗ Improved prompt doesn't include iteration number")
            return False
        
        print("✓ Prompt engineering works correctly")
        print(f"  Original length: {len(original)} chars")
        print(f"  Improved length: {len(improved)} chars")
        
        return True
    except Exception as e:
        print(f"✗ Prompt engineering test failed: {e}")
        return False

def test_file_paths():
    """Test that all required files exist."""
    print("\nTesting file structure...")
    
    base_dir = Path(__file__).parent
    required_files = [
        "scripts/generate_schematic_ai.py",
        "scripts/generate_schematic.py",
        "SKILL.md",
        "README.md"
    ]
    
    all_exist = True
    for file_path in required_files:
        full_path = base_dir / file_path
        if full_path.exists():
            print(f"✓ {file_path}")
        else:
            print(f"✗ Missing: {file_path}")
            all_exist = False
    
    return all_exist

def main():
    """Run all tests."""
    print("="*60)
    print("Scientific Schematics AI Generation - Verification Tests")
    print("="*60)
    
    tests = [
        ("File Structure", test_file_paths),
        ("Imports", test_imports),
        ("Class Structure", test_class_structure),
        ("Error Handling", test_error_handling),
        ("Wrapper Script", test_wrapper_script),
        ("Prompt Engineering", test_prompt_engineering),
    ]
    
    results = []
    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append((test_name, result))
        except Exception as e:
            print(f"\n✗ Test '{test_name}' crashed: {e}")
            results.append((test_name, False))
    
    # Summary
    print("\n" + "="*60)
    print("Test Summary")
    print("="*60)
    
    passed = sum(1 for _, result in results if result)
    total = len(results)
    
    for test_name, result in results:
        status = "✓ PASS" if result else "✗ FAIL"
        print(f"{status}: {test_name}")
    
    print(f"\nTotal: {passed}/{total} tests passed")
    
    if passed == total:
        print("\n✓ All tests passed! Implementation verified.")
        print("\nNext steps:")
        print("1. Set OPENROUTER_API_KEY environment variable")
        print("2. Test with actual API call:")
        print("   python scripts/generate_schematic.py 'test diagram' -o test.png")
        return 0
    else:
        print(f"\n✗ {total - passed} test(s) failed. Please review errors above.")
        return 1

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