gradio-ui/test_export_with_visualization.py

"""
Test export utilities with visualization data
"""

import sys
import os
from pathlib import Path

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from models import (
    ConstructionPlan, PlanMetadata, ExecutiveSummary, RiskData, RiskSummary,
    HazardData, SeismicHazards, VolcanicHazards, HydroHazards, HazardDetail,
    LocationInfo, Coordinates, FacilityInfo, Metadata, Recommendations,
    CostData, CostEstimate, FacilityData, ExportFormats, VisualizationData
)
from export_utils import export_to_json, export_to_pdf, _construction_plan_to_dict


def create_test_construction_plan_with_visualization():
    """Create a test construction plan with visualization data"""
    
    # Create minimal hazard data
    hazard_detail = HazardDetail(
        status="moderate",
        description="Test hazard",
        severity="moderate"
    )
    
    seismic = SeismicHazards(
        active_fault=hazard_detail,
        ground_shaking=hazard_detail,
        liquefaction=hazard_detail,
        tsunami=hazard_detail,
        earthquake_induced_landslide=hazard_detail,
        fissure=hazard_detail,
        ground_rupture=hazard_detail
    )
    
    volcanic = VolcanicHazards(
        active_volcano=hazard_detail,
        potentially_active_volcano=hazard_detail,
        inactive_volcano=hazard_detail,
        ashfall=hazard_detail,
        pyroclastic_flow=hazard_detail,
        lahar=hazard_detail,
        lava=hazard_detail,
        ballistic_projectile=hazard_detail,
        base_surge=hazard_detail,
        volcanic_tsunami=hazard_detail
    )
    
    hydro = HydroHazards(
        flood=hazard_detail,
        rain_induced_landslide=hazard_detail,
        storm_surge=hazard_detail,
        severe_winds=hazard_detail
    )
    
    hazards = HazardData(seismic=seismic, volcanic=volcanic, hydrometeorological=hydro)
    
    summary = RiskSummary(
        overall_risk_level="MODERATE",
        total_hazards_assessed=20,
        high_risk_count=2,
        moderate_risk_count=5,
        critical_hazards=[]
    )
    
    location = LocationInfo(
        name="Test Location",
        coordinates=Coordinates(latitude=14.5995, longitude=120.9842),
        administrative_area="Test Region"
    )
    
    facilities_info = FacilityInfo(schools=[], hospitals=[], road_networks=[])
    
    metadata_risk = Metadata(
        timestamp="2024-01-01T00:00:00",
        source="Test",
        cache_status="test",
        ttl=3600
    )
    
    risk_data = RiskData(
        success=True,
        summary=summary,
        location=location,
        hazards=hazards,
        facilities=facilities_info,
        metadata=metadata_risk
    )
    
    # Create visualization data
    visualization = VisualizationData(
        image_base64="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==",
        prompt_used="Test architectural visualization prompt",
        model_version="gemini-2.5-flash-image",
        generation_timestamp="2024-01-01T00:00:00",
        image_format="PNG",
        resolution="1024x1024",
        features_included=["Reinforced foundation", "Seismic bracing", "Elevated structure"]
    )
    
    # Create construction plan
    metadata = PlanMetadata(
        generated_at="2024-01-01T00:00:00",
        building_type="residential_single_family",
        building_area=100.0,
        location=location,
        coordinates=Coordinates(latitude=14.5995, longitude=120.9842)
    )
    
    executive_summary = ExecutiveSummary(
        overall_risk="MODERATE - 20 hazards assessed",
        critical_concerns=["Test concern"],
        key_recommendations=["Test recommendation"],
        building_specific_notes=["Architectural visualization available showing disaster-resistant features"]
    )
    
    recommendations = Recommendations()
    
    cost_data = CostData(
        total_estimate=CostEstimate(low=100000, mid=150000, high=200000, currency="PHP"),
        market_conditions="Test conditions",
        last_updated="2024-01-01T00:00:00"
    )
    
    facility_data = FacilityData()
    
    export_formats = ExportFormats()
    
    construction_plan = ConstructionPlan(
        metadata=metadata,
        executive_summary=executive_summary,
        risk_assessment=risk_data,
        construction_recommendations=recommendations,
        material_costs=cost_data,
        critical_facilities=facility_data,
        export_formats=export_formats,
        visualization=visualization
    )
    
    return construction_plan


def test_json_export_with_visualization():
    """Test JSON export includes visualization data"""
    print("\n=== Testing JSON Export with Visualization ===")
    
    plan = create_test_construction_plan_with_visualization()
    
    # Convert to dict
    plan_dict = _construction_plan_to_dict(plan)
    
    # Check visualization is included
    assert 'visualization' in plan_dict, "Visualization not in plan dict"
    assert plan_dict['visualization'] is not None, "Visualization is None"
    assert 'image_base64' in plan_dict['visualization'], "image_base64 not in visualization"
    assert 'features_included' in plan_dict['visualization'], "features_included not in visualization"
    
    print("✅ JSON export includes visualization data:")
    print(f"   - Image format: {plan_dict['visualization']['image_format']}")
    print(f"   - Resolution: {plan_dict['visualization']['resolution']}")
    print(f"   - Features: {len(plan_dict['visualization']['features_included'])}")
    
    # Test actual file export
    output_path = "/tmp/test_construction_plan_with_viz.json"
    result_path = export_to_json(plan, output_path)
    
    # Verify file exists
    assert os.path.exists(result_path), f"JSON file not created at {result_path}"
    
    # Read and verify content
    import json
    with open(result_path, 'r') as f:
        exported_data = json.load(f)
    
    assert 'visualization' in exported_data, "Visualization not in exported JSON"
    
    print(f"✅ JSON file exported successfully to {result_path}")
    
    # Cleanup
    os.remove(result_path)
    
    return True


def test_pdf_export_with_visualization():
    """Test PDF export includes visualization image"""
    print("\n=== Testing PDF Export with Visualization ===")
    
    plan = create_test_construction_plan_with_visualization()
    
    # Test PDF generation (will create HTML if weasyprint not available)
    output_path = "/tmp/test_construction_plan_with_viz.pdf"
    
    try:
        result_path = export_to_pdf(plan, output_path)
        
        # Verify file exists
        assert os.path.exists(result_path), f"PDF/HTML file not created at {result_path}"
        
        # Read content and check for visualization section
        with open(result_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # Check for visualization section
        assert '🎨 Architectural Visualization' in content, "Visualization section not in PDF"
        assert 'data:image/png;base64' in content, "Base64 image not embedded in PDF"
        assert 'Disaster-Resistant Features Shown' in content, "Features list not in PDF"
        assert 'Reinforced foundation' in content, "Feature details not in PDF"
        
        print(f"✅ PDF/HTML export includes visualization:")
        print(f"   - Visualization section present")
        print(f"   - Image embedded as base64")
        print(f"   - Features list included")
        print(f"   - File created at {result_path}")
        
        # Cleanup
        os.remove(result_path)
        
        return True
        
    except Exception as e:
        print(f"❌ PDF export failed: {e}")
        return False


def test_export_without_visualization():
    """Test exports work correctly when visualization is None"""
    print("\n=== Testing Export without Visualization ===")
    
    plan = create_test_construction_plan_with_visualization()
    plan.visualization = None
    
    # Test JSON export
    plan_dict = _construction_plan_to_dict(plan)
    assert plan_dict['visualization'] is None, "Visualization should be None"
    
    # Test PDF export
    output_path = "/tmp/test_construction_plan_no_viz.pdf"
    result_path = export_to_pdf(plan, output_path)
    
    with open(result_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Visualization section should not be present
    assert '🎨 Architectural Visualization' not in content, "Visualization section should not be present"
    
    print("✅ Exports work correctly without visualization:")
    print("   - JSON export handles None visualization")
    print("   - PDF export skips visualization section")
    
    # Cleanup
    os.remove(result_path)
    
    return True


def main():
    """Run all tests"""
    print("=" * 60)
    print("EXPORT UTILITIES WITH VISUALIZATION TEST SUITE")
    print("=" * 60)
    
    results = []
    
    try:
        results.append(("JSON Export with Visualization", test_json_export_with_visualization()))
    except Exception as e:
        print(f"❌ JSON export test failed: {e}")
        results.append(("JSON Export with Visualization", False))
    
    try:
        results.append(("PDF Export with Visualization", test_pdf_export_with_visualization()))
    except Exception as e:
        print(f"❌ PDF export test failed: {e}")
        results.append(("PDF Export with Visualization", False))
    
    try:
        results.append(("Export without Visualization", test_export_without_visualization()))
    except Exception as e:
        print(f"❌ Export without visualization test failed: {e}")
        results.append(("Export without Visualization", False))
    
    # Print summary
    print("\n" + "=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)
    
    for test_name, passed in results:
        status = "✅ PASS" if passed else "❌ FAIL"
        print(f"{status}: {test_name}")
    
    passed_count = sum(1 for _, passed in results if passed)
    total_count = len(results)
    
    print(f"\nTotal: {passed_count}/{total_count} test suites passed")
    
    if passed_count == total_count:
        print("\n✅ All tests passed!")
        return 0
    else:
        print(f"\n❌ {total_count - passed_count} test(s) failed")
        return 1


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