109 lines
2.8 KiB
Python
109 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Simple test for the consolidated main.py
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import io
|
|
from PIL import Image
|
|
import base64
|
|
|
|
# Add current directory to PYTHONPATH
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, current_dir)
|
|
|
|
def create_test_image():
|
|
"""Create a simple test image"""
|
|
img = Image.new('RGB', (2048, 2048), color='lightblue')
|
|
from PIL import ImageDraw
|
|
draw = ImageDraw.Draw(img)
|
|
draw.rectangle([500, 500, 1500, 1500], fill='red')
|
|
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format='PNG')
|
|
return buffer.getvalue()
|
|
|
|
def test_preview_function():
|
|
"""Test the preview image creation"""
|
|
# Import the function from main.py
|
|
from main import create_preview_image_b64, get_image_info
|
|
|
|
# Create test image
|
|
test_png = create_test_image()
|
|
print(f"Created test PNG: {len(test_png)} bytes")
|
|
|
|
# Get image info
|
|
info = get_image_info(test_png)
|
|
print(f"Image info: {info}")
|
|
|
|
# Create preview
|
|
preview_b64 = create_preview_image_b64(test_png, target_size=512, quality=85)
|
|
|
|
if preview_b64:
|
|
print(f"✅ Preview created: {len(preview_b64)} chars")
|
|
|
|
# Save preview to verify
|
|
preview_bytes = base64.b64decode(preview_b64)
|
|
with open('test_preview.jpg', 'wb') as f:
|
|
f.write(preview_bytes)
|
|
|
|
ratio = (len(preview_bytes) / len(test_png)) * 100
|
|
print(f"Compression: {ratio:.1f}% (original: {len(test_png)}, preview: {len(preview_bytes)})")
|
|
print("Preview saved as test_preview.jpg")
|
|
return True
|
|
else:
|
|
print("❌ Failed to create preview")
|
|
return False
|
|
|
|
def test_imports():
|
|
"""Test all necessary imports"""
|
|
try:
|
|
from main import (
|
|
Imagen4MCPServer,
|
|
Imagen4ToolHandlers,
|
|
ImageGenerationResult,
|
|
get_tools
|
|
)
|
|
print("✅ All classes imported successfully")
|
|
|
|
# Test tool definitions
|
|
tools = get_tools()
|
|
print(f"✅ {len(tools)} tools defined")
|
|
|
|
return True
|
|
except Exception as e:
|
|
print(f"❌ Import error: {e}")
|
|
return False
|
|
|
|
def main():
|
|
"""Main test function"""
|
|
print("=== Testing Consolidated main.py ===\n")
|
|
|
|
success = True
|
|
|
|
# Test imports
|
|
if not test_imports():
|
|
success = False
|
|
|
|
print()
|
|
|
|
# Test preview function
|
|
if not test_preview_function():
|
|
success = False
|
|
|
|
print(f"\n=== Test Results ===")
|
|
if success:
|
|
print("✅ All tests passed!")
|
|
print("\nTo run the server:")
|
|
print("1. Make sure .env is configured")
|
|
print("2. Run: python main.py")
|
|
print("3. Or use: run.bat")
|
|
else:
|
|
print("❌ Some tests failed")
|
|
|
|
return 0 if success else 1
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|