diff --git a/rocketpy/__init__.py b/rocketpy/__init__.py
index 8db8431e4..10404b619 100644
--- a/rocketpy/__init__.py
+++ b/rocketpy/__init__.py
@@ -1,3 +1,4 @@
+from .control import _Controller
 from .environment import Environment, EnvironmentAnalysis
 from .mathutils import (
     Function,
@@ -22,7 +23,7 @@
     TankGeometry,
     UllageBasedTank,
 )
-from .control import _Controller
+from .plots.compare import Compare, CompareFlights
 from .rocket import (
     AeroSurface,
     AirBrakes,
@@ -37,4 +38,3 @@
     TrapezoidalFins,
 )
 from .simulation import Flight
-from .plots.compare import Compare, CompareFlights
diff --git a/tests/conftest.py b/tests/conftest.py
index 8b9ed9fef..f1accc14e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1279,30 +1279,23 @@ def controller_function(
         z = state[2]
         vz = state[5]
         previous_vz = state_history[-1][5]
-        if time > 3.9:
-            if z < 1500:
-                air_brakes.set_deployment_level(0)
+        if time < 3.9:
+            return None
+        if z < 1500:
+            air_brakes.set_deployment_level(0)
+        else:
+            new_deployment_level = (
+                air_brakes.deployment_level + 0.1 * vz + 0.01 * previous_vz**2
+            )
+            if new_deployment_level > air_brakes.deployment_level + 0.2 / sampling_rate:
+                new_deployment_level = air_brakes.deployment_level + 0.2 / sampling_rate
+            elif (
+                new_deployment_level < air_brakes.deployment_level - 0.2 / sampling_rate
+            ):
+                new_deployment_level = air_brakes.deployment_level - 0.2 / sampling_rate
             else:
-                new_deployment_level = (
-                    air_brakes.deployment_level + 0.1 * vz + 0.01 * previous_vz**2
-                )
-                if (
-                    new_deployment_level
-                    > air_brakes.deployment_level + 0.2 / sampling_rate
-                ):
-                    new_deployment_level = (
-                        air_brakes.deployment_level + 0.2 / sampling_rate
-                    )
-                elif (
-                    new_deployment_level
-                    < air_brakes.deployment_level - 0.2 / sampling_rate
-                ):
-                    new_deployment_level = (
-                        air_brakes.deployment_level - 0.2 / sampling_rate
-                    )
-                else:
-                    new_deployment_level = air_brakes.deployment_level
-                air_brakes.set_deployment_level(new_deployment_level)
+                new_deployment_level = air_brakes.deployment_level
+            air_brakes.set_deployment_level(new_deployment_level)
 
     return controller_function
 
diff --git a/tests/test_flight.py b/tests/test_flight.py
index 0498dc1e3..a2e42281c 100644
--- a/tests/test_flight.py
+++ b/tests/test_flight.py
@@ -311,8 +311,8 @@ def test_air_brakes_flight(mock_show, flight_calisto_air_brakes):
     """
     test_flight = flight_calisto_air_brakes
     air_brakes = test_flight.rocket.air_brakes[0]
-    assert air_brakes.plots.all() == None
-    assert air_brakes.prints.all() == None
+    assert air_brakes.plots.all() is None
+    assert air_brakes.prints.all() is None
 
 
 @patch("matplotlib.pyplot.show")