refactored tests

Author: Gui-FernandesBR

tests/unit/motors/test_point_mass_motor.py

@@ -1,33 +1,52 @@
 import pytest
+
 from rocketpy.motors.point_mass_motor import PointMassMotor
 
+
 def test_init_required_args():
-    m = PointMassMotor(thrust_source=10, dry_mass=1.0, propellant_initial_mass=0.5, burn_time=1.2)
+    m = PointMassMotor(
+        thrust_source=10, dry_mass=1.0, propellant_initial_mass=0.5, burn_time=1.2
+    )
     assert isinstance(m, PointMassMotor)
     assert m.dry_mass == 1.0
     assert m.propellant_initial_mass == 0.5
+    assert m.burn_time == 1.2
+
 
 def test_missing_required_args_raises():
+    # TODO: in the future we would like to capture specific RocketPy Exceptions
     with pytest.raises(ValueError):
         PointMassMotor(10, 1.0, None, 1)
     with pytest.raises(ValueError):
         PointMassMotor(10, 1.0, 1.2)
     with pytest.raises(TypeError):
         PointMassMotor([], 1.0, 1.2, 1)
 
+
 def test_exhaustvelocity_and_totalmassflowrate():
-    m = PointMassMotor(thrust_source=10, dry_mass=1.0, propellant_initial_mass=1.0, burn_time=2.0)
+    m = PointMassMotor(
+        thrust_source=10, dry_mass=1.0, propellant_initial_mass=1.0, burn_time=2.0
+    )
     ev_func = m.exhaustvelocity()
-    assert hasattr(ev_func, 'getValue')
+    assert hasattr(ev_func, "getValue")
     tmf_func = m.totalmassflowrate
-    assert hasattr(tmf_func, 'getValue')
+    assert hasattr(tmf_func, "getValue")
+
 
 def test_zero_inertias():
-    m = PointMassMotor(thrust_source=10, dry_mass=1.0, propellant_initial_mass=1.0, burn_time=2.0)
+    m = PointMassMotor(
+        thrust_source=10, dry_mass=1.0, propellant_initial_mass=1.0, burn_time=2.0
+    )
     assert m.propellantI11().getValue(0) == 0
     assert m.propellantI22().getValue(0) == 0
     assert m.propellantI33().getValue(0) == 0
 
+
 def test_callable_thrust():
-    m = PointMassMotor(thrust_source=lambda t: 100*t, dry_mass=2, propellant_initial_mass=2, burn_time=4)
+    m = PointMassMotor(
+        thrust_source=lambda t: 100 * t,
+        dry_mass=2,
+        propellant_initial_mass=2,
+        burn_time=4,
+    )
     assert m.thrust_source(0.5) == 50

tests/unit/rocket/test_point_mass_rocket.py

@@ -1,5 +1,6 @@
-from rocketpy.rocket.rocket import PointMassRocket
 from rocketpy.motors.point_mass_motor import PointMassMotor
+from rocketpy.rocket.rocket import PointMassRocket
+
 
 def test_point_mass_rocket_basic_properties():
     motor = PointMassMotor(10, 1.0, 0.5, 1.0)
@@ -8,14 +9,15 @@ def test_point_mass_rocket_basic_properties():
         mass=2.0,
         center_of_mass_without_motor=0.1,
         power_off_drag=0.7,
-        power_on_drag=0.8
+        power_on_drag=0.8,
     )
     rocket.addmotor(motor, 0)
     assert rocket.radius == 0.05
     assert rocket.mass == 2.0
     assert rocket.motor is motor
     assert rocket.dryI11 == 0.0  # 3-DOF: inertias are forced zero
 
+
 def test_structural_and_total_mass():
     """Test structural and total mass properties of point mass rocket."""
     motor = PointMassMotor(10, 1.0, 1.1, 2.0)
@@ -24,15 +26,16 @@ def test_structural_and_total_mass():
         mass=2.5,
         center_of_mass_without_motor=0,
         power_off_drag=0.3,
-        power_on_drag=0.4
+        power_on_drag=0.4,
     )
     rocket.addmotor(motor, 0)
-    
+
     # Test that structural mass and total mass are calculated correctly
     assert rocket.mass == 2.5
     expected_total_mass = rocket.mass + motor.propellant_initial_mass
     assert abs(rocket.total_mass(0) - expected_total_mass) < 1e-6
 
+
 def test_add_motor_overwrites():
     """Test that adding a motor overwrites the previous motor."""
     motor1 = PointMassMotor(10, 1, 1.1, 2.0)

tests/unit/simulation/test_flight_3dof.py

@@ -1,43 +1,90 @@
-import pytest
 import numpy as np
-from rocketpy.simulation.flight import Flight
-from rocketpy.rocket.rocket import PointMassRocket
+import pytest
+
 from rocketpy.motors.point_mass_motor import PointMassMotor
+from rocketpy.rocket.point_mass_rocket import PointMassRocket
+from rocketpy.simulation.flight import Flight
+
+
+@pytest.fixture
+def point_mass_motor():
+    """Simple PointMassMotor for 3-DOF tests.
+
+    Returns
+    -------
+    rocketpy.PointMassMotor
+    """
+    return PointMassMotor(
+        thrust_source=10,
+        dry_mass=1.0,
+        propellant_initial_mass=0.5,
+        burn_time=2.2,
+    )
 
-class DummyEnv:
-    # Minimal stub, adapt with real Environment in your RocketPy setup
-    windvelocityx = windvelocityy = speedofsound = pressure = density = dynamicviscosity = lambda self: 0
-    gravity = lambda self: 9.81
-    elevation = 0
-    def __getattr__(self, name):
-        return lambda *a, **k: 0
-
-def make_simple_3dof_components():
-    env = DummyEnv()
-    motor = PointMassMotor(10, dry_mass=1.0, propellant_initial_mass=0.5, burn_time=2.2)
-    rocket = PointMassRocket(0.05, 2.0, 0.1, 0.5, 0.6)
-    rocket.addmotor(motor, 0)
-    return env, rocket
-
-def test_3dof_simulation_mode_autoset():
-    env, rocket = make_simple_3dof_components()
-    flight = Flight(rocket=rocket, environment=env, rail_length=1, simulation_mode="3 DOF")
+
+@pytest.fixture
+def point_mass_rocket(point_mass_motor):
+    """Simple PointMassRocket for 3-DOF tests.
+
+    Returns
+    -------
+    rocketpy.PointMassRocket
+    """
+    rocket = PointMassRocket(
+        radius=0.05,
+        mass=2.0,
+        center_of_mass_without_motor=0.1,
+        power_off_drag=0.5,
+        power_on_drag=0.6,
+    )
+    rocket.add_motor(point_mass_motor, position=0)
+    return rocket
+
+
+def test_3dof_simulation_mode_autoset(example_plain_env, point_mass_rocket):
+    """Test that simulation mode is correctly set to 3 DOF."""
+    flight = Flight(
+        rocket=point_mass_rocket,
+        environment=example_plain_env,
+        rail_length=1,
+        simulation_mode="3 DOF",
+    )
     assert flight.simulation_mode == "3 DOF"
 
-def test_3dof_simulation_mode_warning(monkeypatch):
-    env, rocket = make_simple_3dof_components()
+
+def test_3dof_simulation_mode_warning(
+    monkeypatch, example_plain_env, point_mass_rocket
+):
+    """Test that a warning is issued when 6 DOF is requested with PointMassRocket."""
     monkeypatch.setattr("warnings.warn", lambda *a, **k: None)
-    f = Flight(rocket=rocket, environment=env, rail_length=1, simulation_mode="6 DOF")
-    assert f.simulation_mode == "3 DOF"
+    flight = Flight(
+        rocket=point_mass_rocket,
+        environment=example_plain_env,
+        rail_length=1,
+        simulation_mode="6 DOF",
+    )
+    assert flight.simulation_mode == "3 DOF"
 
-def test_3dof_equations_of_motion_functions():
-    env, rocket = make_simple_3dof_components()
-    flight = Flight(rocket=rocket, environment=env, rail_length=1, simulation_mode="3 DOF")
-    u = [0]*13  # Proper size for generalized state for 3/6 DOF probably
+
+def test_3dof_equations_of_motion_functions(example_plain_env, point_mass_rocket):
+    """Test that 3-DOF equations of motion return valid results."""
+    flight = Flight(
+        rocket=point_mass_rocket,
+        environment=example_plain_env,
+        rail_length=1,
+        simulation_mode="3 DOF",
+    )
+    u = [0] * 13  # Generalized state vector size
     result = flight.udotgeneralized3dof(0, u)
-    assert isinstance(result, list) or isinstance(result, np.ndarray)
+    assert isinstance(result, (list, np.ndarray))
+
 
-def test_invalid_simulation_mode():
-    env, rocket = make_simple_3dof_components()
+def test_invalid_simulation_mode(example_plain_env, point_mass_rocket):
+    """Test that invalid simulation mode raises ValueError."""
     with pytest.raises(ValueError):
-        Flight(rocket=rocket, environment=env, rail_length=1, simulation_mode="2 DOF")
+        Flight(
+            rocket=point_mass_rocket,
+            environment=example_plain_env,
+            rail_length=1,
+            simulation_mode="2 DOF",
+        )