TST: fix tests for new drag implementation

Author: MateusStano

tests/acceptance/test_bella_lui_rocket.py

@@ -141,7 +141,7 @@ def drogue_trigger(p, h, y):
     )
 
     # Define aerodynamic drag coefficients
-    BellaLui.power_off_drag = Function(
+    power_off_drag_by_mach = Function(
         [
             (0.01, 0.51),
             (0.02, 0.46),
@@ -156,7 +156,7 @@ def drogue_trigger(p, h, y):
         "linear",
         "constant",
     )
-    BellaLui.power_on_drag = Function(
+    power_on_drag_by_mach = Function(
         [
             (0.01, 0.51),
             (0.02, 0.46),
@@ -171,8 +171,50 @@ def drogue_trigger(p, h, y):
         "linear",
         "constant",
     )
-    BellaLui.power_off_drag *= parameters.get("power_off_drag")[0]
-    BellaLui.power_on_drag *= parameters.get("power_on_drag")[0]
+    BellaLui.power_off_drag_7d = Function(
+        lambda alpha,
+        beta,
+        mach,
+        reynolds,
+        pitch_rate,
+        yaw_rate,
+        roll_rate: power_off_drag_by_mach.get_value_opt(mach),
+        [
+            "alpha",
+            "beta",
+            "mach",
+            "reynolds",
+            "pitch_rate",
+            "yaw_rate",
+            "roll_rate",
+        ],
+        ["Drag Coefficient with Power Off"],
+        interpolation="linear",
+        extrapolation="constant",
+    )
+    BellaLui.power_on_drag_7d = Function(
+        lambda alpha,
+        beta,
+        mach,
+        reynolds,
+        pitch_rate,
+        yaw_rate,
+        roll_rate: power_on_drag_by_mach.get_value_opt(mach),
+        [
+            "alpha",
+            "beta",
+            "mach",
+            "reynolds",
+            "pitch_rate",
+            "yaw_rate",
+            "roll_rate",
+        ],
+        ["Drag Coefficient with Power On"],
+        interpolation="linear",
+        extrapolation="constant",
+    )
+    BellaLui.power_off_drag_7d *= parameters.get("power_off_drag")[0]
+    BellaLui.power_on_drag_7d *= parameters.get("power_on_drag")[0]
 
     # Flight
     test_flight = Flight(

tests/integration/simulation/test_flight_3dof.py

@@ -200,6 +200,72 @@ def test_weathercock_coeff_default(flight_3dof):
     assert flight_3dof.weathercock_coeff == 0.0
 
 
+def test_point_mass_rocket_3dof_uses_7d_drag_inputs(example_plain_env, point_mass_motor):
+    """Ensure PointMassRocket uses the 7D drag interface in 3-DOF dynamics."""
+
+    def drag_7d(alpha, beta, mach, reynolds, pitch_rate, yaw_rate, roll_rate):
+        return (
+            0.2
+            + 0.01 * abs(alpha)
+            + 0.01 * abs(beta)
+            + 1e-7 * reynolds
+            + 0.001 * (abs(pitch_rate) + abs(yaw_rate) + abs(roll_rate))
+            + 0.01 * mach
+        )
+
+    rocket = PointMassRocket(
+        radius=0.05,
+        mass=2.0,
+        center_of_mass_without_motor=0.1,
+        power_off_drag=drag_7d,
+        power_on_drag=drag_7d,
+    )
+    rocket.add_motor(point_mass_motor, position=0)
+
+    flight = Flight(
+        rocket=rocket,
+        environment=example_plain_env,
+        rail_length=1,
+        simulation_mode="3 DOF",
+    )
+
+    t = 10.0
+    u = [0, 0, 100, 50, 5, 0, 1, 0, 0, 0, 0.3, -0.2, 0.1]
+    u_dot = flight.u_dot_generalized_3dof(t, u)
+
+    z = u[2]
+    vx, vy, vz = u[3], u[4], u[5]
+    omega1, omega2, omega3 = u[10], u[11], u[12]
+
+    rho = flight.env.density.get_value_opt(z)
+    dynamic_viscosity = flight.env.dynamic_viscosity.get_value_opt(z)
+    wind_vx = flight.env.wind_velocity_x.get_value_opt(z)
+    wind_vy = flight.env.wind_velocity_y.get_value_opt(z)
+    speed_of_sound = flight.env.speed_of_sound.get_value_opt(z)
+    gravity = flight.env.gravity.get_value_opt(z)
+
+    free_stream_velocity = np.array([wind_vx - vx, wind_vy - vy, -vz])
+    free_stream_speed = np.linalg.norm(free_stream_velocity)
+    mach = free_stream_speed / speed_of_sound
+
+    stream_velocity_body = free_stream_velocity
+    aerodynamic_stream_velocity = -stream_velocity_body
+    alpha = np.arctan2(aerodynamic_stream_velocity[1], aerodynamic_stream_velocity[2])
+    beta = np.arctan2(aerodynamic_stream_velocity[0], aerodynamic_stream_velocity[2])
+
+    reynolds = (
+        rho * free_stream_speed * (2 * rocket.radius) / dynamic_viscosity
+        if dynamic_viscosity > 0
+        else 0
+    )
+
+    cd_expected = drag_7d(alpha, beta, mach, reynolds, omega1, omega2, omega3)
+    r3_expected = -0.5 * rho * free_stream_speed**2 * rocket.area * cd_expected
+    az_expected = (r3_expected - rocket.total_mass.get_value_opt(t) * gravity) / rocket.total_mass.get_value_opt(t)
+
+    assert u_dot[5] == pytest.approx(az_expected)
+
+
 def test_weathercock_zero_gives_fixed_attitude(flight_weathercock_zero):
     """Tests that weathercock_coeff=0 results in fixed attitude (no quaternion change).
     When weathercock_coeff is 0, the quaternion derivatives should be zero,

tests/unit/rocket/aero_surface/test_generic_surfaces.py

@@ -74,6 +74,32 @@ def test_valid_initialization_from_csv(filename_valid_coeff):
     )
 
 
+def test_csv_independent_variables_accept_any_order(tmp_path):
+    """Checks if GenericSurface correctly maps CSV columns by header names,
+    regardless of independent variable column order."""
+    filename = tmp_path / "valid_coefficients_shuffled_order.csv"
+    filename.write_text(
+        "mach,alpha,cL\n0,0,0\n0,1,10\n2,0,2\n2,1,12\n",
+        encoding="utf-8",
+    )
+
+    generic_surface = GenericSurface(
+        reference_area=REFERENCE_AREA,
+        reference_length=REFERENCE_LENGTH,
+        coefficients={"cL": str(filename)},
+    )
+
+    closure = generic_surface.cL.source.__closure__
+    csv_function = next(
+        cell.cell_contents
+        for cell in closure
+        if isinstance(cell.cell_contents, Function)
+    )
+
+    assert generic_surface.cL(1, 0, 2, 0, 0, 0, 0) == pytest.approx(12)
+    assert csv_function.get_interpolation_method() == "regular_grid"
+
+
 def test_compute_forces_and_moments():
     """Checks if there are not logical errors in
     compute forces and moments"""

tests/unit/rocket/test_rocket.py

@@ -1,4 +1,5 @@
 import warnings
+from itertools import product
 from unittest.mock import patch
 
 import numpy as np
@@ -689,3 +690,146 @@ def test_coordinate_system_orientation(
     static_margin_nose_to_tail = rocket_nose_to_tail.static_margin
 
     assert np.array_equal(static_margin_tail_to_nose, static_margin_nose_to_tail)
+
+
+def test_drag_csv_header_order_independent_for_multivariable_input(tmp_path):
+    """Ensure drag CSV independent-variable columns are interpreted by name.
+
+    This test checks that swapping the order of header-defined variables
+    (mach and reynolds) yields equivalent drag interpolation results.
+    """
+
+    ordered_csv = tmp_path / "drag_mach_reynolds.csv"
+    ordered_csv.write_text(
+        "mach,reynolds,cd\n0.5,0.1,0.6\n1.0,0.1,1.1\n0.5,0.2,0.7\n1.0,0.2,1.2\n",
+        encoding="utf-8",
+    )
+
+    swapped_csv = tmp_path / "drag_reynolds_mach.csv"
+    swapped_csv.write_text(
+        "reynolds,mach,cd\n0.1,0.5,0.6\n0.1,1.0,1.1\n0.2,0.5,0.7\n0.2,1.0,1.2\n",
+        encoding="utf-8",
+    )
+
+    rocket_ordered = Rocket(
+        radius=0.05,
+        mass=1.0,
+        inertia=(1.0, 1.0, 1.0),
+        power_off_drag=str(ordered_csv),
+        power_on_drag=str(ordered_csv),
+        center_of_mass_without_motor=0.0,
+    )
+
+    rocket_swapped = Rocket(
+        radius=0.05,
+        mass=1.0,
+        inertia=(1.0, 1.0, 1.0),
+        power_off_drag=str(swapped_csv),
+        power_on_drag=str(swapped_csv),
+        center_of_mass_without_motor=0.0,
+    )
+
+    drag_ordered = rocket_ordered.power_off_drag_7d(0, 0, 0.8, 0.15, 0, 0, 0)
+    drag_swapped = rocket_swapped.power_off_drag_7d(0, 0, 0.8, 0.15, 0, 0, 0)
+
+    ordered_closure = rocket_ordered.power_off_drag_7d.source.__closure__
+    swapped_closure = rocket_swapped.power_off_drag_7d.source.__closure__
+    ordered_csv_function = next(
+        cell.cell_contents
+        for cell in ordered_closure
+        if isinstance(cell.cell_contents, Function)
+    )
+    swapped_csv_function = next(
+        cell.cell_contents
+        for cell in swapped_closure
+        if isinstance(cell.cell_contents, Function)
+    )
+
+    assert drag_ordered == pytest.approx(0.95)
+    assert drag_swapped == pytest.approx(0.95)
+    assert drag_swapped == pytest.approx(drag_ordered)
+    assert ordered_csv_function.get_interpolation_method() == "regular_grid"
+    assert swapped_csv_function.get_interpolation_method() == "regular_grid"
+
+
+def test_drag_input_types_supported_for_power_on_and_power_off(tmp_path):
+    """Ensure drag input processing accepts all supported input types.
+
+    This test validates that both ``power_off_drag`` and ``power_on_drag``
+    accept and correctly evaluate all supported input categories.
+    """
+    query = (1.0, 0.0, 0.8, 0.15, 0.0, 0.0, 0.0)
+
+    csv_drag = tmp_path / "drag_mach_reynolds.csv"
+    csv_drag.write_text(
+        "mach,reynolds,cd\n0.5,0.1,0.6\n1.0,0.1,1.1\n0.5,0.2,0.7\n1.0,0.2,1.2\n",
+        encoding="utf-8",
+    )
+
+    txt_drag = tmp_path / "drag_curve.txt"
+    txt_drag.write_text("0.0,0.2\n1.0,0.4\n", encoding="utf-8")
+
+    function_1d = Function(
+        lambda mach: 0.2 + mach,
+        inputs=["mach"],
+        outputs=["cd"],
+        interpolation="linear",
+    )
+    function_7d = Function(
+        lambda alpha, beta, mach, reynolds, pitch_rate, yaw_rate, roll_rate: mach
+        + reynolds,
+        inputs=[
+            "alpha",
+            "beta",
+            "mach",
+            "reynolds",
+            "pitch_rate",
+            "yaw_rate",
+            "roll_rate",
+        ],
+        outputs=["cd"],
+        interpolation="linear",
+    )
+
+    drag_7d_table = [
+        (*coords, float(sum(coords))) for coords in product((0.0, 1.0), repeat=7)
+    ]
+    drag_7d_query = (1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0)
+
+    test_cases = [
+        ("int", 1, query, 1.0),
+        ("float", 0.37, query, 0.37),
+        ("csv_path", str(csv_drag), query, 0.95),
+        ("txt_path", str(txt_drag), query, 0.36),
+        ("function_1d", function_1d, query, 1.0),
+        ("function_7d", function_7d, query, 0.95),
+        (
+            "callable_1d",
+            lambda mach: 0.2 + mach,
+            query,
+            1.0,
+        ),
+        (
+            "callable_7d",
+            lambda alpha, beta, mach, reynolds, pitch_rate, yaw_rate, roll_rate: mach
+            + reynolds,
+            query,
+            0.95,
+        ),
+        ("list_pairs", [[0.0, 0.2], [1.0, 0.4]], query, 0.36),
+        ("tuple_pairs", ((0.0, 0.2), (1.0, 0.4)), query, 0.36),
+        ("list_7d_entries", drag_7d_table, drag_7d_query, 4.0),
+    ]
+
+    for _, drag_input, query_point, expected in test_cases:
+        rocket = Rocket(
+            radius=0.05,
+            mass=1.0,
+            inertia=(1.0, 1.0, 1.0),
+            power_off_drag=drag_input,
+            power_on_drag=drag_input,
+            center_of_mass_without_motor=0.0,
+        )
+
+        assert rocket.power_off_drag_7d(*query_point) == pytest.approx(expected)
+        assert rocket.power_on_drag_7d(*query_point) == pytest.approx(expected)