test_orbital_dynamics.py

"""Test suite for OrbitalDynamics module.

Tests high-fidelity orbital dynamics with all perturbations:
- Two-body gravity
- J2 oblateness
- Third-body acceleration
- Solar radiation pressure (with shadowing)
- Atmospheric drag
- Relativistic correction
"""
import pytest
import math
import sys
from pathlib import Path
from unittest.mock import Mock, MagicMock

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent / "orbital-mechanics"))

from orbital_mech.orbital_dynamics import OrbitalDynamics
from models.core import StateDerivative, SpacecraftState, EnvironmentState


class TestOrbitalDynamicsBasic:
    """Test basic orbital dynamics functionality."""
    
    @pytest.fixture
    def mock_spacecraft_state(self):
        """Create a mock SpacecraftState for LEO."""
        state = Mock(spec=SpacecraftState)
        state.position = (6.878137e6, 0, 0)  # 500 km altitude
        state.velocity = (0, 7613, 0)  # Circular orbit velocity
        state.mass = 1000.0  # kg
        state.sail_area = 0.0
        state.sail_normal = (1, 0, 0)
        return state
    
    @pytest.fixture
    def mock_environment_state(self):
        """Create a mock EnvironmentState."""
        return Mock(spec=EnvironmentState)
    
    def test_initialize_with_default_config(self):
        """Test initialization with default Earth config."""
        module = OrbitalDynamics()
        module.initialize({})
        
        assert module.is_valid
        assert module.mu > 0
        assert module.validate()
    
    def test_initialize_with_earth_config(self):
        """Test initialization with Earth-specific config."""
        config = {
            "mu": 3.986004418e14,
            "central_body": "Earth",
            "equatorial_radius": 6.378137e6,
            "J2": 1.081874e-3
        }
        module = OrbitalDynamics()
        module.initialize(config)
        
        assert module.is_valid
        assert module.central_body == "Earth"
        assert module.mu == 3.986004418e14
    
    def test_initialize_with_moon_config(self):
        """Test initialization with Moon config."""
        config = {
            "mu": 4.9048695e12,
            "central_body": "Moon",
            "equatorial_radius": 1.7374e6
        }
        module = OrbitalDynamics()
        module.initialize(config)
        
        assert module.is_valid
        assert module.central_body == "Moon"
    
    def test_compute_derivatives_returns_state_derivative(self, mock_spacecraft_state, mock_environment_state):
        """Test that compute_derivatives returns StateDerivative."""
        module = OrbitalDynamics()
        module.initialize({})
        
        result = module.compute_derivatives(mock_spacecraft_state, mock_environment_state, 0.0)
        
        assert isinstance(result, StateDerivative)
        assert hasattr(result, 'acceleration')
        assert hasattr(result, 'velocity_derivative')
    
    def test_compute_derivatives_never_raises(self, mock_spacecraft_state, mock_environment_state):
        """Test that compute_derivatives never raises exceptions."""
        module = OrbitalDynamics()
        module.initialize({})
        
        try:
            result = module.compute_derivatives(mock_spacecraft_state, mock_environment_state, 0.0)
            assert result is not None
        except Exception as e:
            pytest.fail(f"compute_derivatives raised {e}")
    
    def test_compute_outputs_returns_dict(self, mock_spacecraft_state, mock_environment_state):
        """Test that compute_outputs returns a dict."""
        module = OrbitalDynamics()
        module.initialize({})
        
        result = module.compute_outputs(mock_spacecraft_state, mock_environment_state, 0.0)
        
        assert isinstance(result, dict)
        assert "r_magnitude" in result
        assert "v_magnitude" in result
    
    def test_validate_returns_boolean(self):
        """Test that validate returns boolean."""
        module = OrbitalDynamics()
        module.initialize({})
        
        result = module.validate()
        
        assert isinstance(result, bool)
        assert result is True


class TestOrbitalDynamicsPhysics:
    """Test physical correctness of orbital dynamics."""
    
    @pytest.fixture
    def module(self):
        """Create initialized OrbitalDynamics module."""
        module = OrbitalDynamics()
        module.initialize({
            "mu": 3.986004418e14,
            "central_body": "Earth"
        })
        return module
    
    def test_gravitational_acceleration_magnitude(self, module):
        """Test gravitational acceleration is correct order of magnitude."""
        state = Mock(spec=SpacecraftState)
        state.position = (6.878137e6, 0, 0)  # 500 km altitude
        state.velocity = (0, 0, 0)
        state.mass = 1000.0
        
        env = Mock(spec=EnvironmentState)
        
        deriv = module.compute_derivatives(state, env, 0.0)
        accel_mag = math.sqrt(sum(x**2 for x in deriv.acceleration))
        
        # At 500 km altitude, gravity should be ~8.5 m/s^2
        assert 8.0 < accel_mag < 9.5
    
    def test_circular_orbit_velocity(self, module):
        """Test circular orbit at reference altitude."""
        # Circular orbit at 500 km: v = sqrt(mu/r)
        r = 6.878137e6
        v_circ = math.sqrt(module.mu / r)
        
        state = Mock(spec=SpacecraftState)
        state.position = (r, 0, 0)
        state.velocity = (0, v_circ, 0)
        state.mass = 1000.0
        
        env = Mock(spec=EnvironmentState)
        
        deriv = module.compute_derivatives(state, env, 0.0)
        accel_mag = math.sqrt(sum(x**2 for x in deriv.acceleration))
        
        # Centripetal acceleration should equal gravitational at circular orbit
        a_gravity = module.mu / (r**2)
        assert abs(accel_mag - a_gravity) < 0.01
    
    def test_energy_conservation_two_body(self, module):
        """Test orbital energy is conserved in two-body case."""
        r = 6.878137e6
        v_circ = math.sqrt(module.mu / r)
        
        state = Mock(spec=SpacecraftState)
        state.position = (r, 0, 0)
        state.velocity = (0, v_circ, 0)
        state.mass = 1000.0
        
        env = Mock(spec=EnvironmentState)
        
        # Compute energy at this state
        specific_energy = 0.5 * v_circ**2 - module.mu / r
        
        # For circular orbit, specific energy = -mu/(2*a)
        expected_energy = -module.mu / (2 * r)
        
        # Should match (within numerical precision)
        assert abs(specific_energy - expected_energy) < abs(expected_energy) * 1e-6


class TestOrbitalDynamicsPerturbations:
    """Test perturbation acceleration calculations."""
    
    def test_j2_perturbation_enabled(self):
        """Test J2 perturbation is computed when enabled."""
        module = OrbitalDynamics()
        module.initialize({
            "mu": 3.986004418e14,
            "central_body": "Earth",
            "equatorial_radius": 6.378137e6,
            "J2": 1.081874e-3,
            "enable_J2": True
        })
        
        state = Mock(spec=SpacecraftState)
        state.position = (7.378137e6, 0, 0)  # 1000 km altitude
        state.velocity = (0, 7.5e3, 0)
        state.mass = 1000.0
        
        env = Mock(spec=EnvironmentState)
        
        deriv = module.compute_derivatives(state, env, 0.0)
        
        # J2 component should be non-zero
        j2_mag = math.sqrt(sum(x**2 for x in module.last_components["j2"]))
        assert j2_mag > 0
    
    def test_j2_perturbation_disabled(self):
        """Test J2 perturbation is not computed when disabled."""
        module = OrbitalDynamics()
        module.initialize({
            "mu": 3.986004418e14,
            "enable_J2": False
        })
        
        state = Mock(spec=SpacecraftState)
        state.position = (7.378137e6, 0, 0)
        state.velocity = (0, 7.5e3, 0)
        state.mass = 1000.0
        
        env = Mock(spec=EnvironmentState)
        
        deriv = module.compute_derivatives(state, env, 0.0)
        
        # J2 component should be zero
        j2_mag = math.sqrt(sum(x**2 for x in module.last_components["j2"]))
        assert j2_mag == 0


class TestOrbitalDynamicsInterface:
    """Test PhysicsModule interface compliance."""
    
    def test_get_required_state_fields(self):
        """Test get_required_state_fields returns correct fields."""
        module = OrbitalDynamics()
        module.initialize({})
        
        required = module.get_required_state_fields()
        
        assert isinstance(required, list)
        assert "position" in required
        assert "velocity" in required
        assert "mass" in required
    
    def test_get_updated_state_fields(self):
        """Test get_updated_state_fields returns correct fields."""
        module = OrbitalDynamics()
        module.initialize({})
        
        updated = module.get_updated_state_fields()
        
        assert isinstance(updated, list)
        assert "acceleration" in updated
    
    def test_get_module_type(self):
        """Test get_module_type returns valid type."""
        module = OrbitalDynamics()
        module.initialize({})
        
        module_type = module.get_module_type()
        
        assert module_type is not None
        assert "ORBITAL" in str(module_type).upper()
    
    def test_diagnostics_returns_dict(self):
        """Test diagnostics returns complete dict."""
        module = OrbitalDynamics()
        module.initialize({})
        
        diags = module.diagnostics()
        
        assert isinstance(diags, dict)
        assert "is_valid" in diags
        assert "mu" in diags
        assert "last_error" in diags
    
    def test_get_name(self):
        """Test get_name returns appropriate string."""
        module = OrbitalDynamics()
        module.initialize({})
        
        name = module.get_name()
        
        assert isinstance(name, str)
        assert "OrbitalDynamics" in name
        assert "Earth" in name


class TestOrbitalDynamicsErrorHandling:
    """Test error handling."""
    
    def test_invalid_mu_raises_error_on_init(self):
        """Test that invalid mu causes validation failure."""
        module = OrbitalDynamics()
        module.initialize({"mu": -1000.0})
        
        assert not module.is_valid
        assert module.last_error is not None
    
    def test_invalid_radius_raises_error_on_init(self):
        """Test that invalid radius causes validation failure."""
        module = OrbitalDynamics()
        module.initialize({
            "mu": 3.986004418e14,
            "equatorial_radius": 0.0
        })
        
        assert not module.is_valid
    
    def test_compute_with_invalid_module_returns_zeros(self):
        """Test that compute with invalid module returns zero acceleration."""
        module = OrbitalDynamics()
        module.initialize({"mu": -1.0})  # Invalid
        
        state = Mock(spec=SpacecraftState)
        state.position = (1e7, 0, 0)
        state.velocity = (0, 7000, 0)
        state.mass = 1000.0
        
        env = Mock(spec=EnvironmentState)
        
        deriv = module.compute_derivatives(state, env, 0.0)
        
        # Should return zero acceleration
        assert deriv.acceleration == (0, 0, 0)


class TestOrbitalDynamicsDeterminism:
    """Test that module behaves deterministically."""
    
    def test_same_input_produces_same_output(self):
        """Test deterministic behavior."""
        module = OrbitalDynamics()
        module.initialize({})
        
        state = Mock(spec=SpacecraftState)
        state.position = (6.878137e6, 0, 0)
        state.velocity = (0, 7613, 0)
        state.mass = 1000.0
        
        env = Mock(spec=EnvironmentState)
        
        # Compute twice with same input
        result1 = module.compute_derivatives(state, env, 0.0)
        result2 = module.compute_derivatives(state, env, 0.0)
        
        # Results should be identical
        assert result1.acceleration == result2.acceleration


class TestOrbitalDynamicsConsistency:
    """Test consistency across multiple calls."""
    
    def test_diagnostics_include_acceleration_components(self):
        """Test that diagnostics track all acceleration components."""
        module = OrbitalDynamics()
        module.initialize({
            "enable_J2": True,
            "enable_third_body": False
        })
        
        state = Mock(spec=SpacecraftState)
        state.position = (7.378137e6, 0, 0)
        state.velocity = (0, 7500, 0)
        state.mass = 1000.0
        
        env = Mock(spec=EnvironmentState)
        
        deriv = module.compute_derivatives(state, env, 0.0)
        outputs = module.compute_outputs(state, env, 0.0)
        
        # Should have acceleration component magnitudes in outputs
        assert "accel_gravitational" in outputs
        assert "accel_j2" in outputs
    
    def test_perturbations_additive(self):
        """Test that perturbations add correctly."""
        # Test with only gravity
        module1 = OrbitalDynamics()
        module1.initialize({"enable_J2": False})
        
        # Test with gravity + J2
        module2 = OrbitalDynamics()
        module2.initialize({
            "enable_J2": True,
            "equatorial_radius": 6.378137e6,
            "J2": 1.081874e-3
        })
        
        state = Mock(spec=SpacecraftState)
        state.position = (7.378137e6, 0, 0)
        state.velocity = (0, 7500, 0)
        state.mass = 1000.0
        
        env = Mock(spec=EnvironmentState)
        
        result1 = module1.compute_derivatives(state, env, 0.0)
        result2 = module2.compute_derivatives(state, env, 0.0)
        
        # Result2 acceleration should be different (larger) than result1
        mag1 = math.sqrt(sum(x**2 for x in result1.acceleration))
        mag2 = math.sqrt(sum(x**2 for x in result2.acceleration))
        
        # J2 perturbation should increase total acceleration slightly
        assert mag2 > mag1


if __name__ == "__main__":
    pytest.main([__file__, "-v"])