test_wrappers_simple.py

"""Simple validation test for physics module wrappers.

Quick smoke test to ensure wrapped modules work correctly.
"""
import sys
import os
import math

# Add physics packages to path
WORKSPACE_ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(WORKSPACE_ROOT, 'core-math'))
sys.path.insert(0, os.path.join(WORKSPACE_ROOT, 'orbital-mechanics'))
sys.path.insert(0, os.path.join(WORKSPACE_ROOT, 'sail-physics'))
sys.path.insert(0, WORKSPACE_ROOT)

from models import SpacecraftState, PhysicsInput
from models.wrappers import TwoBodyModule, SolarRadiationPressureModule


def test_two_body():
    """Verify TwoBodyModule works."""
    print("\n" + "="*70)
    print("Test 1: TwoBodyModule (2-body gravity)")
    print("="*70)
    
    from models.core import EnvironmentState
    
    module = TwoBodyModule(mu=1.327e20)
    
    # Test at Earth's distance
    state = SpacecraftState(
        position=(1.496e11, 0, 0),
        velocity=(0, 29.78e3, 0),
        attitude_quaternion=(1, 0, 0, 0),
        angular_velocity=(0, 0, 0),
        mass=260.0,
        power_available=500.0,
        thermal_state={},
        structural_modes=()
    )
    environment = EnvironmentState(
        solar_flux=1361.0,
        magnetic_field=(0, 0, 0),
        solar_wind_density=5e6,
        radiation_index=2.0
    )
    inp = PhysicsInput(state=state, environment=environment, t_epoch=0, params={})
    out = module.compute(inp)
    
    print(f"[OK] Position: {state.position[0]:.2e} m (1 AU)")
    print(f"[OK] Acceleration: {out.state_derivative.acceleration[0]:.6e} m/s² (radial)")
    print(f"[OK] Diagnostics: {list(out.diagnostics.keys())}")
    print(f"[OK] Diagnostics: {out.diagnostics}")
    print(f"[OK] Valid: {out.valid}")
    
    if not out.valid:
        print(f"[FAIL] Error: {out.diagnostics.get('error', 'Unknown error')}")
    
    assert out.valid, "Output should be valid"
    assert out.state_derivative.acceleration[0] < 0, "Should point toward Sun"
    print("[OK] PASSED")


def test_srp():
    """Verify SolarRadiationPressureModule works."""
    print("\n" + "="*70)
    print("Test 2: SolarRadiationPressureModule (SRP)")
    print("="*70)
    
    from models.core import EnvironmentState
    
    module = SolarRadiationPressureModule(
        sail_area_m2=196.0,
        reflectivity=0.85,
        mass_kg=260.0
    )
    
    state = SpacecraftState(
        position=(1.496e11, 0, 0),
        velocity=(0, 30e3, 0),
        attitude_quaternion=(1, 0, 0, 0),
        angular_velocity=(0, 0, 0),
        mass=260.0,
        power_available=500.0,
        thermal_state={},
        structural_modes=()
    )
    environment = EnvironmentState(
        solar_flux=1361.0,
        magnetic_field=(0, 0, 0),
        solar_wind_density=5e6,
        radiation_index=2.0
    )
    inp = PhysicsInput(state=state, environment=environment, t_epoch=0, params={})
    out = module.compute(inp)
    
    print(f"[OK] SRP acceleration: {math.sqrt(sum(x**2 for x in out.state_derivative.acceleration)):.6e} m/s²")
    print(f"[OK] Diagnostics keys: {list(out.diagnostics.keys())}")
    print(f"[OK] Valid: {out.valid}")
    
    assert out.valid, "Output should be valid"
    print("[OK] PASSED")


def test_combined():
    """Verify CombinedOrbitalDynamicsModule works."""
    print("\n" + "="*70)
    print("Test 3: CombinedOrbitalDynamicsModule")
    print("="*70)
    
    from models.wrappers import CombinedOrbitalDynamicsModule
    from models.core import EnvironmentState
    
    module = CombinedOrbitalDynamicsModule(
        mu=1.327e20,
        include_perturbations=False  # Disable perturbations for this test
    )
    
    state = SpacecraftState(
        position=(1.496e11, 0, 0),
        velocity=(0, 29780, 0),
        attitude_quaternion=(1, 0, 0, 0),
        angular_velocity=(0, 0, 0),
        mass=260.0,
        power_available=500.0,
        thermal_state={},
        structural_modes=()
    )
    environment = EnvironmentState(
        solar_flux=1361.0,
        magnetic_field=(0, 0, 0),
        solar_wind_density=5e6,
        radiation_index=2.0
    )
    inp = PhysicsInput(state=state, environment=environment, t_epoch=0, params={})
    out = module.compute(inp)
    
    print(f"[OK] Total acceleration: {math.sqrt(sum(x**2 for x in out.state_derivative.acceleration)):.6e} m/s²")
    print(f"[OK] Radial component: {out.state_derivative.acceleration[0]:.6e} m/s²")
    print(f"[OK] Valid: {out.valid}")
    
    assert out.valid, "Output should be valid"
    assert out.state_derivative.acceleration[0] < 0, "Should point toward Sun"
    print("[OK] PASSED")


def main():
    print("\n" + "="*70)
    print("PHYSICS MODULE WRAPPER VALIDATION")
    print("="*70)
    
    try:
        test_two_body()
        test_srp()
        test_combined()
        
        print("\n" + "="*70)
        print("[OK] ALL TESTS PASSED")
        print("="*70)
        return 0
    
    except Exception as e:
        print(f"\n[FAIL] TEST FAILED: {e}")
        import traceback
        traceback.print_exc()
        return 1


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