test_kmz.py

#!/usr/bin/env python3
"""
Test script for KMZ export functionality
"""
import numpy as np
import zipfile
import xml.etree.ElementTree as ET
import os

def create_test_radar_kml(lat2d, lon2d, z2d, forecast_hour=6, min_dbz=1.0):
    """Test version of create_radar_kml function."""
    try:
        # Create KML structure
        kml = ET.Element('kml', xmlns="http://www.opengis.net/kml/2.2")
        document = ET.SubElement(kml, 'Document')
        
        # Add document info
        name = ET.SubElement(document, 'name')
        name.text = f"Test RAP Radar Forecast +{forecast_hour}h (min {min_dbz} dBZ)"
        
        description = ET.SubElement(document, 'description')
        description.text = f"Test RAP Composite Reflectivity forecast for +{forecast_hour} hours"
        
        # Add ground overlay for bounds
        ground_overlay = ET.SubElement(document, 'GroundOverlay')
        overlay_name = ET.SubElement(ground_overlay, 'name')
        overlay_name.text = f"RAP Test Radar Grid"
        
        # Create boundaries
        lat_box = ET.SubElement(ground_overlay, 'LatLonBox')
        north = ET.SubElement(lat_box, 'north')
        south = ET.SubElement(lat_box, 'south') 
        east = ET.SubElement(lat_box, 'east')
        west = ET.SubElement(lat_box, 'west')
        
        north.text = str(float(np.nanmax(lat2d)))
        south.text = str(float(np.nanmin(lat2d)))
        east.text = str(float(np.nanmax(lon2d)))
        west.text = str(float(np.nanmin(lon2d)))
        
        # Add a few test placemarks
        for i in range(0, min(5, lat2d.shape[0] - 1)):
            for j in range(0, min(5, lat2d.shape[1] - 1)):
                value = z2d[i, j]
                
                if np.isnan(value) or value < min_dbz:
                    continue
                    
                # Create placemark for this cell
                placemark = ET.SubElement(document, 'Placemark')
                pm_name = ET.SubElement(placemark, 'name')
                pm_name.text = f"Test {value:.1f} dBZ"
                
                pm_desc = ET.SubElement(placemark, 'description')
                pm_desc.text = f"Test radar reflectivity: {value:.1f} dBZ"
                
                # Create point
                point = ET.SubElement(placemark, 'Point')
                coordinates = ET.SubElement(point, 'coordinates')
                coordinates.text = f"{lon2d[i, j]},{lat2d[i, j]},0"
        
        # Convert to string
        rough_string = ET.tostring(kml, 'unicode')
        return rough_string
        
    except Exception as e:
        print(f"Test KML creation error: {e}")
        return f"""<?xml version="1.0" encoding="UTF-8"?>
<kml xmlns="http://www.opengis.net/kml/2.2">
  <Document>
    <name>Test RAP Radar Export Error</name>
    <description>Error creating test KML: {str(e)}</description>
  </Document>
</kml>"""

def test_kmz_export():
    """Test the KMZ export functionality with simulated data."""
    print("=== Testing KMZ Export Functionality ===")
    
    try:
        # Create test data (simplified CONUS region)
        ny, nx = 20, 30
        lat_min, lat_max = 35, 45
        lon_min, lon_max = -105, -85
        
        # Create test grid
        lat_1d = np.linspace(lat_min, lat_max, ny)
        lon_1d = np.linspace(lon_min, lon_max, nx)
        lon2d, lat2d = np.meshgrid(lon_1d, lat_1d)
        
        # Create test radar data with some realistic patterns
        z2d = np.random.rand(ny, nx) * 30 + 5  # 5-35 dBZ range
        
        # Add some NaN values to test filtering
        z2d[z2d < 8] = np.nan
        
        print(f"Test grid: {ny}x{nx}")
        print(f"Lat range: {lat2d.min():.2f} to {lat2d.max():.2f}")
        print(f"Lon range: {lon2d.min():.2f} to {lon2d.max():.2f}")
        print(f"Radar values: {np.nanmin(z2d):.1f} to {np.nanmax(z2d):.1f} dBZ")
        
        # Test KML creation
        print("\n--- Testing KML creation ---")
        kml_content = create_test_radar_kml(lat2d, lon2d, z2d, forecast_hour=6, min_dbz=1.0)
        print(f"KML content length: {len(kml_content)} characters")
        
        # Test KMZ creation
        print("\n--- Testing KMZ creation ---")
        os.makedirs('test_exports', exist_ok=True)
        kmz_path = "test_exports/test_rap_radar.kmz"
        
        with zipfile.ZipFile(kmz_path, 'w', zipfile.ZIP_DEFLATED) as kmz:
            kmz.writestr('doc.kml', kml_content)
            
        print(f"KMZ file created: {kmz_path}")
        print(f"File size: {os.path.getsize(kmz_path)} bytes")
        
        # Verify KMZ content
        print("\n--- Verifying KMZ content ---")
        with zipfile.ZipFile(kmz_path, 'r') as kmz:
            files = kmz.namelist()
            print(f"Files in KMZ: {files}")
            
            if 'doc.kml' in files:
                kml_data = kmz.read('doc.kml').decode('utf-8')
                print(f"KML data length: {len(kml_data)} characters")
                
                # Check for key elements
                if '<kml' in kml_data and '<Document>' in kml_data:
                    print("✓ Valid KML structure found")
                else:
                    print("✗ Invalid KML structure")
            else:
                print("✗ doc.kml not found in KMZ")
        
        print(f"\n=== KMZ export test completed successfully ===")
        print(f"Generated file: {os.path.abspath(kmz_path)}")
        print("You can open this file in Google Earth to test the visualization.")
        
        return True
        
    except Exception as e:
        print(f"KMZ export test failed: {e}")
        return False

if __name__ == "__main__":
    test_kmz_export()