test_environment.py

import time
from datetime import date, datetime, timedelta, timezone
from unittest.mock import patch
import numpy as np
import pytest


@pytest.mark.parametrize(
    "lat, lon, theoretical_elevation",
    [
        (48.858844, 2.294351, 34),  # The Eiffel Tower
        (32.990254, -106.974998, 1401),  # Spaceport America
    ],
)
def test_set_elevation_open_elevation(
    lat, lon, theoretical_elevation, example_plain_env
):
    example_plain_env.set_location(lat, lon)

    # either successfully gets the elevation or raises RuntimeError
    try:
        example_plain_env.set_elevation(elevation="Open-Elevation")
        assert example_plain_env.elevation == pytest.approx(
            theoretical_elevation, abs=1
        ), "The Open-Elevation API returned an unexpected value for the elevation"
    except RuntimeError:
        pass  # Ignore the error and pass the test


@patch("matplotlib.pyplot.show")
def test_era5_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
    """Tests the Reanalysis model with the ERA5 file. It uses an example file
    available in the data/weather folder of the RocketPy repository.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_spaceport_env : rocketpy.Environment
        Example environment object to be tested.
    """
    example_spaceport_env.set_date((2018, 10, 15, 12))
    example_spaceport_env.set_atmospheric_model(
        type="Reanalysis",
        file="data/weather/SpaceportAmerica_2018_ERA-5.nc",
        dictionary="ECMWF",
    )
    assert example_spaceport_env.all_info() is None


@patch("matplotlib.pyplot.show")
def test_custom_atmosphere(mock_show, example_plain_env):  # pylint: disable=unused-argument
    """Tests the custom atmosphere model in the environment object.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_plain_env : rocketpy.Environment
        Example environment object to be tested.
    """
    example_plain_env.set_atmospheric_model(
        type="custom_atmosphere",
        pressure=None,
        temperature=300,
        wind_u=[(0, 5), (1000, 10)],
        wind_v=[(0, -2), (500, 3), (1600, 2)],
    )
    assert example_plain_env.all_info() is None
    assert abs(example_plain_env.pressure(0) - 101325.0) < 1e-8
    assert abs(example_plain_env.barometric_height(101325.0)) < 1e-2
    assert abs(example_plain_env.wind_velocity_x(0) - 5) < 1e-8
    assert abs(example_plain_env.temperature(100) - 300) < 1e-8


@patch("matplotlib.pyplot.show")
def test_standard_atmosphere(mock_show, example_plain_env):  # pylint: disable=unused-argument
    """Tests the standard atmosphere model in the environment object.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_plain_env : rocketpy.Environment
        Example environment object to be tested.
    """
    example_plain_env.set_atmospheric_model(type="standard_atmosphere")
    assert example_plain_env.info() is None
    assert example_plain_env.all_info() is None
    assert abs(example_plain_env.pressure(0) - 101325.0) < 1e-8
    assert abs(example_plain_env.barometric_height(101325.0)) < 1e-2
    assert example_plain_env.prints.print_earth_details() is None


@patch("matplotlib.pyplot.show")
def test_wind_plots_wrapping_direction(mock_show, example_plain_env):  # pylint: disable=unused-argument
    """Tests that wind direction plots handle 360°→0° wraparound without
    drawing a horizontal line across the graph.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_plain_env : rocketpy.Environment
        Example environment object to be tested.
    """
    # Set a custom atmosphere where wind direction wraps from ~350° to ~10°
    # across the altitude range by choosing wind_u and wind_v to create a
    # direction near 350° at low altitude and ~10° at higher altitude.
    # wind_direction = (180 + atan2(wind_u, wind_v)) % 360
    # For direction ~350°: need atan2(wind_u, wind_v) ≈ 170° → wind_u>0, wind_v<0
    # For direction ~10°:  need atan2(wind_u, wind_v) ≈ -170° → wind_u<0, wind_v<0
    example_plain_env.set_atmospheric_model(
        type="custom_atmosphere",
        pressure=None,
        temperature=300,
        wind_u=[(0, 1), (5000, -1)],  # changes sign across altitude
        wind_v=[(0, -6), (5000, -6)],  # stays negative → heading near 350°/10°
    )
    # Verify that the wind direction actually wraps through 0°/360° in this
    # atmosphere so the test exercises the wraparound code path.
    low_dir = example_plain_env.wind_direction(0)
    high_dir = example_plain_env.wind_direction(5000)
    assert abs(low_dir - high_dir) > 180, (
        "Test setup error: wind direction should cross 0°/360° boundary"
    )
    # Verify that the helper inserts NaN breaks into the direction and altitude
    # arrays at the wraparound point, which is the core of the fix.
    directions = np.array(
        [example_plain_env.wind_direction(i) for i in example_plain_env.plots.grid],
        dtype=float,
    )
    altitudes = np.array(example_plain_env.plots.grid, dtype=float)
    directions_broken, altitudes_broken = (
        example_plain_env.plots._break_direction_wraparound(directions, altitudes)
    )
    assert np.any(np.isnan(directions_broken)), (
        "Expected NaN breaks in direction array at 0°/360° wraparound"
    )
    assert np.any(np.isnan(altitudes_broken)), (
        "Expected NaN breaks in altitude array at 0°/360° wraparound"
    )
    # Verify info() and atmospheric_model() plots complete without error
    assert example_plain_env.info() is None
    assert example_plain_env.plots.atmospheric_model() is None


@pytest.mark.slow
@pytest.mark.parametrize(
    "model_name",
    [
        "ECMWF",
        "GFS",
        "ICON",
        "ICONEU",
    ],
)
def test_windy_atmosphere(example_euroc_env, model_name):
    """Tests the Windy model in the environment object. The test ensures the
    pressure, temperature, and wind profiles are working and giving reasonable
    values. The tolerances may be higher than usual due to the nature of the
    atmospheric uncertainties, but it is ok since we are just testing if the
    method is working.

    Parameters
    ----------
    example_euroc_env : Environment
        Example environment object to be tested. The EuRoC launch site is used
        to test the ICONEU model, which only works in Europe.
    model_name : str
        The name of the model to be passed to the set_atmospheric_model() method
        as the "file" parameter.
    """
    example_euroc_env.set_atmospheric_model(type="Windy", file=model_name)
    assert pytest.approx(100000.0, rel=0.1) == example_euroc_env.pressure(100)
    assert 0 + 273 < example_euroc_env.temperature(100) < 40 + 273
    assert abs(example_euroc_env.wind_velocity_x(100)) < 20.0
    assert abs(example_euroc_env.wind_velocity_y(100)) < 20.0


@pytest.mark.slow
@patch("matplotlib.pyplot.show")
def test_gfs_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
    """Tests the Forecast model with the GFS file. It does not test the values,
    instead the test checks if the method runs without errors.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_spaceport_env : rocketpy.Environment
        Example environment object to be tested.
    """
    example_spaceport_env.set_atmospheric_model(type="Forecast", file="GFS")
    assert example_spaceport_env.all_info() is None


@pytest.mark.slow
@patch("matplotlib.pyplot.show")
def test_aigfs_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
    """Tests the Forecast model with the AIGFS file.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_spaceport_env : rocketpy.Environment
        Example environment object to be tested.
    """
    example_spaceport_env.set_atmospheric_model(type="Forecast", file="AIGFS")
    assert example_spaceport_env.all_info() is None


@pytest.mark.slow
@patch("matplotlib.pyplot.show")
def test_hrrr_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
    """Tests the Forecast model with the HRRR file.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_spaceport_env : rocketpy.Environment
        Example environment object to be tested.
    """
    # Sometimes the HRRR latest-model can fail due to not having at least 24
    # hours in the future in the forecast, so we try with 12 hours in the future
    # only.
    example_spaceport_env.set_date(datetime.now() + timedelta(hours=12))
    example_spaceport_env.set_atmospheric_model(type="Forecast", file="HRRR")
    assert example_spaceport_env.all_info() is None


@pytest.mark.slow
@patch("matplotlib.pyplot.show")
def test_nam_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
    """Tests the Forecast model with the NAM file.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_spaceport_env : rocketpy.Environment
        Example environment object to be tested.
    """
    example_spaceport_env.set_atmospheric_model(type="Forecast", file="NAM")
    assert example_spaceport_env.all_info() is None


@pytest.mark.slow
@patch("matplotlib.pyplot.show")
def test_rap_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
    now = datetime.now(timezone.utc)
    example_spaceport_env.set_date((now.year, now.month, now.day, now.hour))
    example_spaceport_env.set_atmospheric_model(type="Forecast", file="RAP")
    assert example_spaceport_env.all_info() is None


@pytest.mark.slow
@patch("matplotlib.pyplot.show")
def test_gefs_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
    """Tests the Ensemble model with the GEFS file.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_spaceport_env : rocketpy.Environment
        Example environment object to be tested.
    """
    with pytest.raises(
        ValueError,
        match="GEFS latest-model shortcut is currently unavailable",
    ):
        example_spaceport_env.set_atmospheric_model(type="Ensemble", file="GEFS")


@pytest.mark.slow
@patch("matplotlib.pyplot.show")
def test_wyoming_sounding_atmosphere(mock_show, example_plain_env):  # pylint: disable=unused-argument
    """Asserts whether the Wyoming sounding model in the environment
    object behaves as expected with respect to some attributes such
    as pressure, barometric_height, wind_velocity and temperature.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_plain_env : rocketpy.Environment
        Example environment object to be tested.
    """

    # TODO:: this should be added to the set_atmospheric_model() method as a
    #        "file" option, instead of receiving the URL as a string.
    url = "http://weather.uwyo.edu/cgi-bin/sounding?region=samer&TYPE=TEXT%3ALIST&YEAR=2019&MONTH=02&FROM=0500&TO=0512&STNM=83779"
    # give it at least 5 times to try to download the file
    for i in range(5):
        try:
            example_plain_env.set_atmospheric_model(type="wyoming_sounding", file=url)
            break
        except Exception:  # pylint: disable=broad-except
            time.sleep(2**i)
    assert example_plain_env.all_info() is None
    assert abs(example_plain_env.pressure(0) - 93600.0) < 1e-8
    assert (
        abs(example_plain_env.barometric_height(example_plain_env.pressure(0)) - 722.0)
        < 1e-8
    )
    assert abs(example_plain_env.wind_velocity_x(0) - -2.9005178894925043) < 1e-8
    assert abs(example_plain_env.temperature(100) - 291.75) < 1e-8


@pytest.mark.slow
@patch("matplotlib.pyplot.show")
def test_hiresw_ensemble_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
    """Tests the Forecast model with the HIRESW file.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_spaceport_env : rocketpy.Environment
        Example environment object to be tested.
    """
    today = date.today()
    date_info = (today.year, today.month, today.day, 12)  # Hour given in UTC time

    example_spaceport_env.set_date(date_info)

    with pytest.raises(
        ValueError,
        match="HIRESW latest-model shortcut is currently unavailable",
    ):
        example_spaceport_env.set_atmospheric_model(
            type="Forecast",
            file="HIRESW",
            dictionary="HIRESW",
        )


@pytest.mark.skip(reason="CMC model is currently not working")
@patch("matplotlib.pyplot.show")
def test_cmc_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
    """Tests the Ensemble model with the CMC file.

    Parameters
    ----------
    mock_show : mock
        Mock object to replace matplotlib.pyplot.show() method.
    example_spaceport_env : rocketpy.Environment
        Example environment object to be tested.
    """
    example_spaceport_env.set_atmospheric_model(type="Ensemble", file="CMC")
    assert example_spaceport_env.all_info() is None


def test_merra2_full_specification_compliance(merra2_file_path, example_plain_env):
    """
    Tests that RocketPy loads a file complying with NASA MERRA-2 file specs.
    """
    # 1. Initialize Environment
    env = example_plain_env
    env.set_date((2023, 6, 20, 12))
    env.set_location(latitude=0, longitude=0)

    # 2. Force standard gravity to a known constant for precise math checking
    env.standard_g = 9.80665

    # 3. Load the Atmospheric Model (Using the file generated above)
    env.set_atmospheric_model(
        type="Reanalysis", file=merra2_file_path, dictionary="MERRA2"
    )

    # 4. Verify Unit Conversion (Energy -> Height)
    # Input: 9806.65 m2/s2
    # Expected: 1000.0 m
    print(f"Calculated Elevation: {env.elevation} m")
    assert abs(env.elevation - 1000.0) < 1e-4, (
        f"Failed to convert PHIS (m2/s2) to meters. Got {env.elevation}, expected 1000.0"
    )

    # 5. Verify Variable Mapping
    assert abs(env.temperature(0) - 300.0) < 1e-6
    assert env.wind_speed(0) > 0