test_io_infrastructure.py

"""Test I/O infrastructure: Configuration, trajectory export, audit logging.

Validates:
- YAML configuration loading
- HDF5 trajectory export
- SQLite audit log export
- Round-trip data integrity
"""

import sys
import logging
from pathlib import Path

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

# Import from local modules using absolute paths
import importlib.util

# Load config_loader
spec = importlib.util.spec_from_file_location("config_loader", 
    str(Path(__file__).parent / "io" / "config_loader.py"))
config_loader = importlib.util.module_from_spec(spec)
spec.loader.exec_module(config_loader)

# Load trajectory_export
spec = importlib.util.spec_from_file_location("trajectory_export",
    str(Path(__file__).parent / "io" / "trajectory_export.py"))
trajectory_export = importlib.util.module_from_spec(spec)
spec.loader.exec_module(trajectory_export)

# Load audit_export
spec = importlib.util.spec_from_file_location("audit_export",
    str(Path(__file__).parent / "io" / "audit_export.py"))
audit_export = importlib.util.module_from_spec(spec)
spec.loader.exec_module(audit_export)

load_config = config_loader.load_config
save_config = config_loader.save_config
export_trajectory_hdf5 = trajectory_export.export_trajectory_hdf5
load_trajectory_hdf5 = trajectory_export.load_trajectory_hdf5
export_audit_log_sqlite = audit_export.export_audit_log_sqlite
AuditDatabase = audit_export.AuditDatabase

# Load io __init__ for export_all
spec = importlib.util.spec_from_file_location("io_init",
    str(Path(__file__).parent / "io" / "__init__.py"))
io_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(io_module)
export_all = io_module.export_all

from models.wrappers import TwoBodyModule, SolarRadiationPressureModule
from api import Mission


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def test_config_loading():
    """Test YAML configuration loading."""
    print("\n" + "="*80)
    print("TEST 1: YAML Configuration Loading")
    print("="*80)
    
    config_file = "missions/solar_sail_2body.yaml"
    config = load_config(config_file)
    
    assert config.name == "solar_sail_2body", "Mission name mismatch"
    assert config.t_end == 604800.0, "Mission duration mismatch"
    assert config.spacecraft.sail_area_m2 == 196.0, "Sail area mismatch"
    assert config.initial_state is not None, "Initial state not loaded"
    assert config.initial_state.r == (7e9, 0, 0), "Initial position mismatch"
    
    print(f"✓ Config loaded: {config.name}")
    print(f"  Duration: {config.t_end} s ({config.t_end / 86400:.1f} days)")
    print(f"  Spacecraft: {config.spacecraft.name}")
    print(f"  Sail area: {config.spacecraft.sail_area_m2} m²")
    print(f"  Initial position: {config.initial_state.r}")
    print(f"  Initial velocity: {config.initial_state.v}")
    
    return config


def test_mission_execution(config):
    """Test mission execution with physics modules."""
    print("\n" + "="*80)
    print("TEST 2: Mission Execution with Physics Modules")
    print("="*80)
    
    # Create physics modules
    orbital = TwoBodyModule(mu=1.327e20)
    srp = SolarRadiationPressureModule(
        area_m2=config.spacecraft.sail_area_m2,
        reflectivity=config.spacecraft.sail_reflectivity,
        mass_kg=config.spacecraft.mass_dry_kg
    )
    
    # Create and run mission
    mission = Mission(config)
    mission.add_physics_module("orbital_mechanics", orbital)
    mission.add_physics_module("sail_srp", srp)
    
    result = mission.run()
    
    print(f"✓ Mission completed: {result.config.name}")
    print(f"  Total steps: {len(result.trajectory)}")
    print(f"  Wall time: {result.elapsed_wall_time_sec:.2f} s")
    print(f"  Steps per second: {len(result.trajectory) / result.elapsed_wall_time_sec:.0f}")
    print(f"  Success: {result.success}")
    
    # Check trajectory
    assert len(result.trajectory) > 0, "Empty trajectory"
    assert result.success, "Mission failed"
    
    first = result.trajectory[0]
    last = result.trajectory[-1]
    
    print(f"  Initial r: {first.state.r[0]:.2e} m")
    print(f"  Final r:   {last.state.r[0]:.2e} m")
    print(f"  Initial v: {first.state.v[1]:.0f} m/s")
    print(f"  Final v:   {last.state.v[1]:.0f} m/s")
    
    return result


def test_trajectory_export(result):
    """Test HDF5 trajectory export and reload."""
    print("\n" + "="*80)
    print("TEST 3: HDF5 Trajectory Export and Reload")
    print("="*80)
    
    output_dir = Path("test_outputs")
    output_dir.mkdir(exist_ok=True)
    h5_file = output_dir / "test_trajectory.h5"
    
    # Export
    export_trajectory_hdf5(result, str(h5_file))
    
    assert h5_file.exists(), "HDF5 file not created"
    file_size_mb = h5_file.stat().st_size / 1e6
    
    print(f"✓ Exported trajectory to {h5_file}")
    print(f"  File size: {file_size_mb:.2f} MB")
    
    # Reload
    data = load_trajectory_hdf5(str(h5_file))
    
    assert "metadata" in data, "Metadata missing"
    assert "state" in data, "State missing"
    assert "time" in data, "Time missing"
    
    metadata = data["metadata"]
    print(f"  Metadata n_steps: {metadata['n_steps']}")
    print(f"  Mission name: {metadata['mission_name']}")
    print(f"  Spacecraft: {metadata['spacecraft_name']}")
    print(f"  Solver: {metadata['solver_type']}")
    
    # Verify state arrays
    state = data["state"]
    assert len(state["r"]) == len(result.trajectory), "Position array size mismatch"
    assert len(state["v"]) == len(result.trajectory), "Velocity array size mismatch"
    assert len(state["q"]) == len(result.trajectory), "Quaternion array size mismatch"
    
    print(f"✓ Reloaded trajectory data")
    print(f"  Position array shape: {state['r'].shape}")
    print(f"  Velocity array shape: {state['v'].shape}")
    print(f"  Attitude array shape: {state['q'].shape}")
    
    return h5_file


def test_audit_export(result):
    """Test SQLite audit trail export and query."""
    print("\n" + "="*80)
    print("TEST 4: SQLite Audit Trail Export and Query")
    print("="*80)
    
    output_dir = Path("test_outputs")
    output_dir.mkdir(exist_ok=True)
    db_file = output_dir / "test_audit.db"
    
    # Export
    export_audit_log_sqlite(result, str(db_file))
    
    assert db_file.exists(), "SQLite database not created"
    file_size_kb = db_file.stat().st_size / 1e3
    
    print(f"✓ Exported audit trail to {db_file}")
    print(f"  File size: {file_size_kb:.2f} KB")
    
    # Query
    db = AuditDatabase(str(db_file))
    db.connect()
    
    try:
        # Query mission metadata
        cursor = db.conn.cursor()
        cursor.execute("SELECT * FROM mission_metadata WHERE id = 1")
        mission_row = cursor.fetchone()
        
        assert mission_row is not None, "Mission metadata not found"
        print(f"✓ Mission metadata found")
        print(f"  Name: {mission_row[1]}")  # mission_name
        print(f"  Seed: {mission_row[2]}")  # config_seed
        print(f"  Steps: {mission_row[5]}")  # n_steps
        print(f"  Wall time: {mission_row[6]:.2f} s")  # wall_time_sec
        print(f"  Success: {mission_row[7]}")  # success
        
        # Query events
        events = db.query_events()
        print(f"✓ Events queried: {len(events)} events")
        
        # Query audit log
        audit_entries = db.query_audit_log(limit=5)
        print(f"✓ Audit log queried: {len(audit_entries)} recent entries")
        if audit_entries:
            for entry in audit_entries[:3]:
                print(f"    {entry['action_type']}: {entry['description']}")
        
    finally:
        db.close()
    
    return db_file


def test_export_all(result):
    """Test unified export function."""
    print("\n" + "="*80)
    print("TEST 5: Unified Export Function")
    print("="*80)
    
    output_dir = Path("test_outputs/unified_export")
    
    from io import export_all
    export_all(result, str(output_dir), prefix=result.config.name)
    
    assert (output_dir / f"{result.config.name}.h5").exists(), "HDF5 file missing"
    assert (output_dir / f"{result.config.name}_audit.db").exists(), "Audit DB missing"
    assert (output_dir / f"{result.config.name}_config.yaml").exists(), "Config YAML missing"
    
    print(f"✓ All files exported to {output_dir}")
    
    h5_size = (output_dir / f"{result.config.name}.h5").stat().st_size / 1e6
    db_size = (output_dir / f"{result.config.name}_audit.db").stat().st_size / 1e3
    yaml_size = (output_dir / f"{result.config.name}_config.yaml").stat().st_size
    
    print(f"  Trajectory (HDF5): {h5_size:.2f} MB")
    print(f"  Audit trail (SQLite): {db_size:.2f} KB")
    print(f"  Configuration (YAML): {yaml_size} bytes")
    
    # Reload configuration
    config_file = output_dir / f"{result.config.name}_config.yaml"
    reloaded_config = load_config(str(config_file))
    
    assert reloaded_config.name == result.config.name, "Config name mismatch"
    assert reloaded_config.t_end == result.config.t_end, "Config t_end mismatch"
    
    print(f"✓ Configuration reloaded and verified")


def main():
    """Run all I/O tests."""
    print("\n" + "="*80)
    print("HELIOSAIL-RX I/O INFRASTRUCTURE TEST SUITE")
    print("="*80)
    
    try:
        # Test 1: Config loading
        config = test_config_loading()
        
        # Test 2: Mission execution
        result = test_mission_execution(config)
        
        # Test 3: HDF5 export
        h5_file = test_trajectory_export(result)
        
        # Test 4: SQLite audit export
        db_file = test_audit_export(result)
        
        # Test 5: Unified export
        test_export_all(result)
        
        print("\n" + "="*80)
        print("ALL TESTS PASSED ✓")
        print("="*80)
        print("\nI/O Infrastructure Summary:")
        print("  ✓ YAML configuration loading")
        print("  ✓ Mission execution with 2 physics modules")
        print("  ✓ HDF5 trajectory export (60,480+ steps)")
        print("  ✓ SQLite audit trail export")
        print("  ✓ Data integrity verification")
        print("  ✓ Round-trip configuration reloading")
        
        return 0
        
    except Exception as e:
        print(f"\n✗ TEST FAILED: {e}")
        import traceback
        traceback.print_exc()
        return 1


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