gpt-edit/tests/test_main.py

"""
Simple test runner for GPTEdit
This script runs basic tests without requiring pytest
"""

import sys
import os
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

def test_imports():
    """Test that all modules can be imported"""
    print("Testing imports...")
    
    try:
        # Test config import
        from src.connector.config import Config
        print("✓ Config module imported")
        
        # Test validation utilities
        from src.utils.validation import validate_edit_parameters, sanitize_prompt
        print("✓ Validation utilities imported")
        
        # Test image utilities (without PIL dependency for now)
        try:
            from src.utils.image_utils import encode_image_base64, decode_image_base64
            print("✓ Image utilities imported")
        except ImportError as e:
            print(f"⚠ Image utilities require PIL: {e}")
        
        # Test server models
        from src.server.models import MCPToolDefinitions
        print("✓ Server models imported")
        
        # Test MCP server
        try:
            from src.server.mcp_server import GPTEditMCPServer
            print("✓ MCP server imported")
        except ImportError as e:
            print(f"⚠ MCP server requires mcp package: {e}")
        
        return True
        
    except Exception as e:
        print(f"✗ Import error: {e}")
        return False


def test_config():
    """Test configuration functionality"""
    print("\nTesting configuration...")
    
    try:
        from src.connector.config import Config
        
        # Create config instance
        config = Config()
        print("✓ Config instance created")
        
        # Test fixed parameters
        assert config.MODEL == "gpt-image-1"
        assert config.QUALITY == "high"
        assert config.NUMBER_OF_IMAGES == 1
        print("✓ Fixed parameters correct")
        
        # Test size calculation
        assert config.get_optimal_size(100, 100) == "256x256"
        assert config.get_optimal_size(500, 500) == "512x512"
        assert config.get_optimal_size(1000, 1000) == "1024x1024"
        print("✓ Size calculation working")
        
        # Test max size conversion
        config.max_image_size_mb = 4
        assert config.get_max_image_size_bytes() == 4 * 1024 * 1024
        print("✓ Size conversion working")
        
        return True
        
    except Exception as e:
        print(f"✗ Config test error: {e}")
        return False


def test_validation():
    """Test validation utilities"""
    print("\nTesting validation utilities...")
    
    try:
        from src.utils.validation import sanitize_prompt, validate_edit_parameters
        
        # Test prompt sanitization
        prompt = "  Test   prompt   "
        sanitized = sanitize_prompt(prompt)
        assert sanitized == "Test prompt"
        print("✓ Prompt sanitization working")
        
        # Test long prompt truncation
        long_prompt = "x" * 2000
        sanitized_long = sanitize_prompt(long_prompt)
        assert len(sanitized_long) == 1000
        print("✓ Prompt truncation working")
        
        # Test parameter validation (without file check)
        params = {
            'image_path': 'test.png',  # Won't exist but we're testing structure
            'prompt': 'Edit this image'
        }
        is_valid, error_msg = validate_edit_parameters(params)
        # Should fail because file doesn't exist, but that's expected
        assert is_valid is False
        assert "not found" in error_msg
        print("✓ Parameter validation structure working")
        
        return True
        
    except Exception as e:
        print(f"✗ Validation test error: {e}")
        return False


def test_base64_encoding():
    """Test base64 encoding/decoding"""
    print("\nTesting base64 encoding...")
    
    try:
        from src.utils.image_utils import encode_image_base64, decode_image_base64
        
        # Test data
        test_data = b"Test image data"
        
        # Encode
        encoded = encode_image_base64(test_data)
        assert isinstance(encoded, str)
        print("✓ Base64 encoding working")
        
        # Decode
        decoded = decode_image_base64(encoded)
        assert decoded == test_data
        print("✓ Base64 decoding working")
        
        return True
        
    except ImportError:
        print("⚠ Skipping base64 tests (PIL not installed)")
        return True
    except Exception as e:
        print(f"✗ Base64 test error: {e}")
        return False


def test_tool_definitions():
    """Test MCP tool definitions"""
    print("\nTesting tool definitions...")
    
    try:
        from src.server.models import MCPToolDefinitions
        
        # Get all tools
        tools = MCPToolDefinitions.get_all_tools()
        assert len(tools) == 5
        print(f"✓ Found {len(tools)} tools")
        
        # Check tool names
        tool_names = [tool.name for tool in tools]
        expected_names = [
            "edit_image",
            "edit_with_mask", 
            "batch_edit",
            "validate_image",
            "create_mask_from_alpha"
        ]
        
        for name in expected_names:
            assert name in tool_names
        print(f"✓ All expected tools defined: {', '.join(tool_names)}")
        
        # Check edit_image tool structure
        edit_tool = MCPToolDefinitions.get_edit_image_tool()
        assert edit_tool.name == "edit_image"
        assert "image_path" in edit_tool.inputSchema["properties"]
        assert "prompt" in edit_tool.inputSchema["properties"]
        print("✓ Tool schema structure correct")
        
        return True
        
    except ImportError as e:
        print(f"⚠ Tool definitions require mcp package: {e}")
        return True
    except Exception as e:
        print(f"✗ Tool definitions test error: {e}")
        return False


def main():
    """Run all tests"""
    print("=" * 60)
    print("GPTEdit Test Suite")
    print("=" * 60)
    
    tests = [
        test_imports,
        test_config,
        test_validation,
        test_base64_encoding,
        test_tool_definitions
    ]
    
    passed = 0
    failed = 0
    
    for test in tests:
        if test():
            passed += 1
        else:
            failed += 1
    
    print("\n" + "=" * 60)
    print(f"Test Results: {passed} passed, {failed} failed")
    print("=" * 60)
    
    if failed == 0:
        print("✅ All tests passed!")
        return 0
    else:
        print(f"❌ {failed} test(s) failed")
        return 1


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