gradio-ui/test_visualization_display.py

"""
Test visualization display component
"""

import sys
import os
import base64
import tempfile

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from components.visualization_display import format_visualization_display


def test_visualization_display_with_data():
    """Test visualization display with valid data"""
    # Create a simple 1x1 PNG image in base64
    # This is a minimal valid PNG (1x1 red pixel)
    png_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
    
    visualization_data = {
        'image_base64': png_base64,
        'prompt_used': 'Test prompt for architectural visualization',
        'model_version': 'gemini-2.5-flash-image',
        'generation_timestamp': '2024-01-01T12:00:00Z',
        'image_format': 'PNG',
        'resolution': '1024x1024',
        'features_included': [
            'Reinforced concrete frame',
            'Seismic bracing',
            'Elevated foundation'
        ]
    }
    
    html, image_path, download_path = format_visualization_display(visualization_data)
    
    # Verify HTML contains expected elements
    assert '🎨 Architectural Visualization' in html
    assert 'gemini-2.5-flash-image' in html
    assert 'Test prompt for architectural visualization' in html
    assert 'Reinforced concrete frame' in html
    assert 'Seismic bracing' in html
    assert 'Elevated foundation' in html
    
    # Verify image path is created
    assert image_path is not None
    assert os.path.exists(image_path)
    assert image_path.endswith('.png')
    
    # Verify download path is created
    assert download_path is not None
    assert os.path.exists(download_path)
    assert download_path.endswith('.png')
    
    # Clean up
    if image_path and os.path.exists(image_path):
        os.remove(image_path)
    if download_path and os.path.exists(download_path):
        os.remove(download_path)
    
    print("✅ Test passed: visualization display with valid data")


def test_visualization_display_no_data():
    """Test visualization display with no data"""
    html, image_path, download_path = format_visualization_display(None)
    
    assert 'No Visualization Available' in html
    assert image_path is None
    assert download_path is None
    
    print("✅ Test passed: visualization display with no data")


def test_visualization_display_empty_dict():
    """Test visualization display with empty dict"""
    html, image_path, download_path = format_visualization_display({})
    
    assert 'No Visualization Available' in html
    assert image_path is None
    assert download_path is None
    
    print("✅ Test passed: visualization display with empty dict")


def test_visualization_display_error():
    """Test visualization display with error"""
    visualization_data = {
        'success': False,
        'error': {
            'code': 'GENERATION_FAILED',
            'message': 'Image generation failed',
            'retry_possible': True
        }
    }
    
    html, image_path, download_path = format_visualization_display(visualization_data)
    
    assert 'Visualization Generation Failed' in html
    assert 'GENERATION_FAILED' in html
    assert 'Image generation failed' in html
    assert image_path is None
    assert download_path is None
    
    print("✅ Test passed: visualization display with error")


def test_visualization_display_minimal_data():
    """Test visualization display with minimal data (no features)"""
    png_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
    
    visualization_data = {
        'image_base64': png_base64,
        'prompt_used': '',
        'model_version': 'test-model',
        'generation_timestamp': '2024-01-01T12:00:00Z',
        'image_format': 'PNG',
        'resolution': '512x512',
        'features_included': []
    }
    
    html, image_path, download_path = format_visualization_display(visualization_data)
    
    assert '🎨 Architectural Visualization' in html
    assert 'test-model' in html
    assert image_path is not None
    assert os.path.exists(image_path)
    assert download_path is not None
    assert os.path.exists(download_path)
    
    # Clean up
    if image_path and os.path.exists(image_path):
        os.remove(image_path)
    if download_path and os.path.exists(download_path):
        os.remove(download_path)
    
    print("✅ Test passed: visualization display with minimal data")


if __name__ == '__main__':
    print("Running visualization display component tests...\n")
    
    test_visualization_display_with_data()
    test_visualization_display_no_data()
    test_visualization_display_empty_dict()
    test_visualization_display_error()
    test_visualization_display_minimal_data()
    
    print("\n✅ All tests passed!")