diff --git a/.gitignore b/.gitignore
index f7acbbb..2eb6b5d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -251,6 +251,6 @@ $RECYCLE.BIN/
 **.html
 
 # Command sequences
-*.fmu
+#*.fmu
 
 # End of https://www.gitignore.io/api/git,linux,macos,python,windows,pycharm,jupyternotebook
diff --git a/rocket_twin/drivers/mission.py b/rocket_twin/drivers/mission.py
index a8ecbe1..5544079 100644
--- a/rocket_twin/drivers/mission.py
+++ b/rocket_twin/drivers/mission.py
@@ -43,9 +43,6 @@ def __init__(
         # Init and stop conditions
         init_fuel = init
         init_flight = {
-            "rocket.flying": True,
-            "controller.w_temp": 0.0,
-            "rocket.stage_1.controller.w_temp": 1.0,
             "rocket.stage_1.tank.fuel.w_out_max": 3.0,
             "g_tank.w_in": 0.0,
         }
diff --git a/rocket_twin/systems/__init__.py b/rocket_twin/systems/__init__.py
index 4506d75..366c376 100644
--- a/rocket_twin/systems/__init__.py
+++ b/rocket_twin/systems/__init__.py
@@ -1,4 +1,9 @@
-from rocket_twin.systems.control import ControllerCoSApp, RocketControllerCoSApp, ControllerFMU, RocketControllerFMU
+from rocket_twin.systems.control import (
+    ControllerCoSApp,
+    ControllerFMU,
+    RocketControllerCoSApp,
+    RocketControllerFMU,
+)
 from rocket_twin.systems.engine import Engine, EngineGeom, EnginePerfo
 from rocket_twin.systems.ground import Ground
 from rocket_twin.systems.physics import Dynamics
diff --git a/rocket_twin/systems/control/__init__.py b/rocket_twin/systems/control/__init__.py
index 19c02b9..1d861bd 100644
--- a/rocket_twin/systems/control/__init__.py
+++ b/rocket_twin/systems/control/__init__.py
@@ -2,6 +2,6 @@
 from rocket_twin.systems.control.rocket_controller_cosapp import RocketControllerCoSApp
 
 from rocket_twin.systems.control.controller_fmu import ControllerFMU  # isort: skip
-from rocket_twin.systems.control.rocket_controller_fmu import RocketControllerFMU # isort: skip
+from rocket_twin.systems.control.rocket_controller_fmu import RocketControllerFMU  # isort: skip
 
 __all__ = ["ControllerCoSApp", "ControllerFMU", "RocketControllerCoSApp", "RocketControllerFMU"]
diff --git a/rocket_twin/systems/control/controller/controller.fmu b/rocket_twin/systems/control/controller/controller.fmu
new file mode 100644
index 0000000..f421ecd
Binary files /dev/null and b/rocket_twin/systems/control/controller/controller.fmu differ
diff --git a/rocket_twin/systems/control/controller_cosapp.py b/rocket_twin/systems/control/controller_cosapp.py
index 6578724..c5addf0 100644
--- a/rocket_twin/systems/control/controller_cosapp.py
+++ b/rocket_twin/systems/control/controller_cosapp.py
@@ -17,8 +17,8 @@ def setup(self):
 
         self.add_inward("is_on", 0, desc="Whether the fuel flow is allowed or not")
 
-        self.add_outward("w", 1., desc="Ratio of command fuel flow to maximum flow", unit="")
+        self.add_outward("w", 1.0, desc="Ratio of command fuel flow to maximum flow", unit="")
 
     def compute(self):
 
-        self.w = self.is_on
\ No newline at end of file
+        self.w = self.is_on
diff --git a/rocket_twin/systems/control/rocket_controller/rocket_controller.fmu b/rocket_twin/systems/control/rocket_controller/rocket_controller.fmu
new file mode 100644
index 0000000..77fb547
Binary files /dev/null and b/rocket_twin/systems/control/rocket_controller/rocket_controller.fmu differ
diff --git a/rocket_twin/systems/control/rocket_controller_cosapp.py b/rocket_twin/systems/control/rocket_controller_cosapp.py
index f344158..a042885 100644
--- a/rocket_twin/systems/control/rocket_controller_cosapp.py
+++ b/rocket_twin/systems/control/rocket_controller_cosapp.py
@@ -1,24 +1,26 @@
 from cosapp.base import System
 
-class RocketControllerCoSApp(System):
 
+class RocketControllerCoSApp(System):
     def setup(self, n_stages):
 
         self.add_inward("n_stages", n_stages, desc="number of stages")
         self.add_inward("stage", 1, desc="Current active stage")
 
         for i in range(1, n_stages + 1):
-            self.add_inward(f"weight_prop_{i}", 0., desc=f"Stage {i} propellant weight", unit='kg')
-            self.add_inward(f"weight_max_{i}", 1., desc=f"Stage {i} maximum propellant weight", unit='kg')
+            self.add_inward(f"weight_prop_{i}", 0.0, desc=f"Stage {i} propellant weight", unit="kg")
+            self.add_inward(
+                f"weight_max_{i}", 1.0, desc=f"Stage {i} maximum propellant weight", unit="kg"
+            )
             self.add_outward(f"is_on_{i}", 0, desc=f"Whether the stage {i} is on or not")
 
-        self.add_inward("time_int", 5., desc="Interval between fueling end and launch", unit='s')
-        self.add_inward("time_lnc", 100000., desc="Launch time", unit='s')
+        self.add_inward("time_int", 5.0, desc="Interval between fueling end and launch", unit="s")
+        self.add_inward("time_lnc", 100000.0, desc="Launch time", unit="s")
 
         self.add_outward("fueling", 1, desc="Whether the rocket is fueling or not")
         self.add_outward("flying", 0, desc="Whether the rocket is flying or not")
 
-        self.add_event("full", trigger = "weight_prop_1 == weight_max_1")
+        self.add_event("full", trigger="weight_prop_1 == weight_max_1")
         self.add_event("launch", trigger="t == time_lnc")
         self.add_event("drop", trigger="weight_prop_1 == 0.")
 
@@ -37,9 +39,9 @@ def transition(self):
             self.is_on_1 = 1
         if self.drop.present:
             if self.stage < self.n_stages:
-                self[f'is_on_{self.stage}'] = 0
+                self[f"is_on_{self.stage}"] = 0
                 self.stage += 1
-                self[f'is_on_{self.stage}'] = 1
+                self[f"is_on_{self.stage}"] = 1
                 self.drop.trigger = f"weight_prop_{self.stage} == 0."
             else:
-                self[f'is_on_{self.stage}'] = 0
+                self[f"is_on_{self.stage}"] = 0
diff --git a/rocket_twin/systems/control/rocket_controller_fmu.py b/rocket_twin/systems/control/rocket_controller_fmu.py
index 84547a4..25136b2 100644
--- a/rocket_twin/systems/control/rocket_controller_fmu.py
+++ b/rocket_twin/systems/control/rocket_controller_fmu.py
@@ -6,8 +6,8 @@
 
 import rocket_twin.systems.control
 
-class RocketControllerFMU(System):
 
+class RocketControllerFMU(System):
     def setup(self, n_stages, model_path, model_name):
 
         self.add_inward("n_stages", n_stages, desc="number of stages")
@@ -18,7 +18,7 @@ def setup(self, n_stages, model_path, model_name):
         self.add_inward("time_lnc", 100000.0, desc="Launch time", unit="")
         self.add_transient("x", der="1")
 
-        pulling = {"flying" : "flying", "fueling" : "fueling", "tl" : "time_lnc", "ti" : "time_var"}
+        pulling = {"flying": "flying", "fueling": "fueling", "tl": "time_lnc", "ti": "time_var"}
 
         for i in range(1, n_stages + 1):
             self.add_outward(f"is_on_{i}", 0, desc=f"Whether the stage {i} is on or not")
@@ -32,7 +32,7 @@ def setup(self, n_stages, model_path, model_name):
             pulling=pulling,
         )
 
-        self.add_event("full", trigger = "weight_prop_1 > 0.9999*weight_max_1")
+        self.add_event("full", trigger="weight_prop_1 > 0.9999*weight_max_1")
         self.add_event("drop", trigger="weight_prop_1 < 0.1")
 
     def compute(self):
@@ -48,7 +48,7 @@ def transition(self):
             else:
                 self.time_lnc = self.time + self.time_int
                 self.stage = 1
-        
+
         if self.drop.present:
             if self.stage < self.n_stages:
                 self.stage += 1
@@ -86,5 +86,3 @@ def create_fmu(self, model_path, model_name):
                 os.remove(filename)
 
         return fmu
-
-
diff --git a/rocket_twin/systems/rocket/rocket.py b/rocket_twin/systems/rocket/rocket.py
index 7479a2a..9b6b8c5 100644
--- a/rocket_twin/systems/rocket/rocket.py
+++ b/rocket_twin/systems/rocket/rocket.py
@@ -1,8 +1,8 @@
 from cosapp.base import System
+from OCC.Core.GProp import GProp_GProps
 
 from rocket_twin.systems import Dynamics, RocketControllerCoSApp
 from rocket_twin.systems.rocket import OCCGeometry, Stage
-from OCC.Core.GProp import GProp_GProps
 
 
 class Rocket(System):
@@ -53,13 +53,23 @@ def setup(self, n_stages=1):
             properties[i - 1] = f"stage_{i}"
             forces[i - 1] = f"thrust_{i}"
 
-        self.add_child(RocketControllerCoSApp("controller", n_stages=n_stages), execution_index=0, pulling=['fueling', 'flying'])
+        self.add_child(
+            RocketControllerCoSApp("controller", n_stages=n_stages),
+            execution_index=0,
+            pulling=["fueling", "flying"],
+        )
         self.add_child(OCCGeometry("geom", shapes=shapes, properties=properties))
         self.add_child(Dynamics("dyn", forces=forces, weights=["weight_rocket"]), pulling=["a"])
 
         for i in range(1, n_stages + 1):
-            self.connect(self.controller.outwards, self[f"stage_{i}"].inwards, {f"is_on_{i}" : "is_on"})
-            self.connect(self[f"stage_{i}"].outwards, self.controller.inwards, {"weight_prop" : f"weight_prop_{i}", "weight_max" : f"weight_max_{i}"})
+            self.connect(
+                self.controller.outwards, self[f"stage_{i}"].inwards, {f"is_on_{i}": "is_on"}
+            )
+            self.connect(
+                self[f"stage_{i}"].outwards,
+                self.controller.inwards,
+                {"weight_prop": f"weight_prop_{i}", "weight_max": f"weight_max_{i}"},
+            )
             self.connect(self[f"stage_{i}"].outwards, self.geom.inwards, {"props": f"stage_{i}"})
             self.connect(self[f"stage_{i}"].outwards, self.dyn.inwards, {"thrust": f"thrust_{i}"})
 
@@ -77,4 +87,3 @@ def transition(self):
                 self.geom[f"stage_{self.stage}"] = GProp_GProps()
                 self.dyn[f"thrust_{self.stage}"] = 0.0
                 self.stage += 1
-
diff --git a/rocket_twin/systems/rocket/stage.py b/rocket_twin/systems/rocket/stage.py
index 8aee2be..e1cecb8 100644
--- a/rocket_twin/systems/rocket/stage.py
+++ b/rocket_twin/systems/rocket/stage.py
@@ -25,7 +25,7 @@ def setup(self, nose=False, wings=False):
         shapes = ["tank_s", "engine_s", "tube_s"]
         properties = ["tank", "engine", "tube"]
 
-        self.add_child(ControllerCoSApp("controller"), pulling=['is_on'])
+        self.add_child(ControllerCoSApp("controller"), pulling=["is_on"])
         self.add_child(Tank("tank"), pulling=["w_in", "weight_max", "weight_prop"])
         self.add_child(Engine("engine"), pulling={"force": "thrust"})
         self.add_child(TubeGeom("tube"))
diff --git a/rocket_twin/systems/station/station.py b/rocket_twin/systems/station/station.py
index 204f01e..03a87da 100644
--- a/rocket_twin/systems/station/station.py
+++ b/rocket_twin/systems/station/station.py
@@ -1,5 +1,4 @@
 from cosapp.base import System
-from OCC.Core.GProp import GProp_GProps
 
 from rocket_twin.systems import ControllerCoSApp, Pipe, Rocket, Tank
 
@@ -27,7 +26,7 @@ def setup(self, n_stages=1):
         self.connect(self.g_tank.outwards, self.pipe.inwards, {"w_out": "w_in"})
         self.connect(self.pipe.outwards, self.rocket.inwards, {"w_out": "w_in_1"})
 
-        self.connect(self.rocket.outwards, self.controller.inwards, {'fueling' : 'is_on'})
+        self.connect(self.rocket.outwards, self.controller.inwards, {"fueling": "is_on"})
         self.connect(self.controller.outwards, self.g_tank.inwards, {"w": "w_command"})
 
         self.g_tank.geom.height = 2.0
diff --git a/rocket_twin/tests/test_controller_cosapp.py b/rocket_twin/tests/test_controller_cosapp.py
index 9adbb10..c06ce40 100644
--- a/rocket_twin/tests/test_controller_cosapp.py
+++ b/rocket_twin/tests/test_controller_cosapp.py
@@ -1,19 +1,19 @@
 import numpy as np
+from cosapp.drivers import NonLinearSolver, RungeKutta
+from cosapp.recorders import DataFrameRecorder
 
 from rocket_twin.systems import Station
-from cosapp.drivers import RungeKutta, NonLinearSolver
-from cosapp.recorders import DataFrameRecorder
 
-class TestControllerCosapp:
 
+class TestControllerCosapp:
     def test_run_once(self):
 
-        sys = Station('sys', n_stages=3)
+        sys = Station("sys", n_stages=3)
         sys.run_once()
 
     def test_control(self):
 
-        sys = Station('sys', n_stages=3)
+        sys = Station("sys", n_stages=3)
 
         init = {
             "g_tank.fuel.weight_p": 20.0,
@@ -21,7 +21,7 @@ def test_control(self):
             "rocket.stage_1.tank.fuel.w_out_max": 1.0,
             "rocket.stage_2.tank.fuel.w_out_max": 1.0,
             "rocket.stage_3.tank.fuel.w_out_max": 1.0,
-            "rocket.controller.time_int" : 5.
+            "rocket.controller.time_int": 5.0,
         }
 
         includes = [
@@ -29,11 +29,11 @@ def test_control(self):
             "rocket.weight_prop_1",
             "rocket.weight_prop_2",
             "rocket.weight_prop_3",
-            "rocket.a"
+            "rocket.a",
         ]
 
         driver = sys.add_driver(RungeKutta("rk", order=4, dt=1))
-        solver = driver.add_child(NonLinearSolver('solver'))
+        solver = driver.add_child(NonLinearSolver("solver"))
         driver.time_interval = (0, 35)
         driver.set_scenario(init=init)
         driver.add_recorder(DataFrameRecorder(includes=includes), period=1.0)
@@ -41,12 +41,17 @@ def test_control(self):
         sys.run_drivers()
 
         data = driver.recorder.export_data()
-        data1 = data.drop(["Section", "Status", "Error code", "rocket.weight_prop_2", "rocket.weight_prop_3"], axis=1)
-        data2 = data.drop(["Section", "Status", "Error code", "rocket.dyn.weight", "rocket.weight_prop_1"], axis=1)
+        data1 = data.drop(
+            ["Section", "Status", "Error code", "rocket.weight_prop_2", "rocket.weight_prop_3"],
+            axis=1,
+        )
+        data2 = data.drop(
+            ["Section", "Status", "Error code", "rocket.dyn.weight", "rocket.weight_prop_1"], axis=1
+        )
         acel = np.asarray(data["rocket.a"])
 
         print(data1)
         print(data2)
 
-        np.testing.assert_allclose(sys.rocket.geom.weight, 4., atol=10 ** (0))
-        np.testing.assert_allclose(acel[-3], 40., atol=0.1)
\ No newline at end of file
+        np.testing.assert_allclose(sys.rocket.geom.weight, 4.0, atol=10 ** (0))
+        np.testing.assert_allclose(acel[-3], 40.0, atol=0.1)
diff --git a/test_controller_fmu.py b/rocket_twin/tests/test_controller_fmu.py
similarity index 59%
rename from test_controller_fmu.py
rename to rocket_twin/tests/test_controller_fmu.py
index 5975fe0..9342250 100644
--- a/test_controller_fmu.py
+++ b/rocket_twin/tests/test_controller_fmu.py
@@ -24,26 +24,30 @@ def test_controller_fmu(self):
         )
         swap_system(
             sys.rocket.controller,
-            RocketControllerFMU("controller", n_stages=n_stages, model_path=model_path_r, model_name=model_name_r),
+            RocketControllerFMU(
+                "controller", n_stages=n_stages, model_path=model_path_r, model_name=model_name_r
+            ),
         )
         for i in range(1, n_stages + 1):
             swap_system(
                 sys.rocket[f"stage_{i}"].controller,
-                ControllerFMU("controller", model_path=model_path, model_name=model_name)
+                ControllerFMU("controller", model_path=model_path, model_name=model_name),
             )
 
         driver = sys.add_driver(RungeKutta(order=4, time_interval=[0, 40], dt=0.1))
-        solver = driver.add_child(NonLinearSolver("solver"))
-        init = {"g_tank.fuel.weight_p": 20.0,
-                "rocket.stage_1.tank.fuel.weight_p": 0.0,
-                "rocket.stage_2.tank.fuel.weight_p": 0.0,
-                "rocket.stage_3.tank.fuel.weight_p": 0.0}
+        driver.add_child(NonLinearSolver("solver"))
+        init = {
+            "g_tank.fuel.weight_p": 20.0,
+            "rocket.stage_1.tank.fuel.weight_p": 0.0,
+            "rocket.stage_2.tank.fuel.weight_p": 0.0,
+            "rocket.stage_3.tank.fuel.weight_p": 0.0,
+        }
         values = {
             "g_tank.fuel.w_out_max": 1.0,
             "rocket.controller.time_int": 5.0,
-            "rocket.stage_1.tank.fuel.w_out_max": 1.,
-            "rocket.stage_2.tank.fuel.w_out_max": 1.,
-            "rocket.stage_3.tank.fuel.w_out_max": 1.,
+            "rocket.stage_1.tank.fuel.w_out_max": 1.0,
+            "rocket.stage_2.tank.fuel.w_out_max": 1.0,
+            "rocket.stage_3.tank.fuel.w_out_max": 1.0,
         }
         driver.set_scenario(init=init, values=values)
         driver.add_recorder(
@@ -53,12 +57,15 @@ def test_controller_fmu(self):
         sys.run_drivers()
         data = driver.recorder.export_data()
         data1 = data.drop(["Section", "Status", "Error code", "rocket.a"], axis=1)
-        data2 = data.drop(["Section", "Status", "Error code", "rocket.geom.weight", "rocket.weight_prop_3"], axis=1)
+        data2 = data.drop(
+            ["Section", "Status", "Error code", "rocket.geom.weight", "rocket.weight_prop_3"],
+            axis=1,
+        )
 
         print(data1)
         print(data2)
 
-        np.testing.assert_allclose(sys.rocket.geom.weight, 4., atol=10 ** (0))
-        np.testing.assert_allclose(sys.rocket.a, 40., rtol=0.1)
-        #np.testing.assert_allclose(sys.g_tank.weight_prop, 5.0, atol=10 ** (0))
-        #np.testing.assert_allclose(sys.rocket.stage_1.tank.weight_prop, 0.0, atol=10 ** (0))
+        np.testing.assert_allclose(sys.rocket.geom.weight, 4.0, atol=10 ** (0))
+        np.testing.assert_allclose(sys.rocket.a, 40.0, rtol=0.1)
+        # np.testing.assert_allclose(sys.g_tank.weight_prop, 5.0, atol=10 ** (0))
+        # np.testing.assert_allclose(sys.rocket.stage_1.tank.weight_prop, 0.0, atol=10 ** (0))
diff --git a/rocket_twin/tests/test_drivers.py b/rocket_twin/tests/test_drivers.py
new file mode 100644
index 0000000..b373c66
--- /dev/null
+++ b/rocket_twin/tests/test_drivers.py
@@ -0,0 +1,70 @@
+import numpy as np
+
+from rocket_twin.drivers.fueling_rocket import FuelingRocket
+from rocket_twin.drivers.vertical_flying_rocket import VerticalFlyingRocket
+from rocket_twin.systems import Station
+
+
+class TestDrivers:
+    """Tests for the FuelingRocket and VerticalFlyingRocket driver."""
+
+    sys = Station("sys")
+
+    def test_fuel(self):
+        dt = 0.1
+
+        init = {
+            "rocket.stage_1.tank.fuel.weight_p": 0.0,
+            "g_tank.fuel.weight_p": 10.0,
+            "g_tank.w_in": 0.0,
+            "g_tank.fuel.w_out_max": 3.0,
+        }
+
+        stop = "rocket.controller.flying == 1."
+
+        includes = ["rocket.a"]
+
+        self.sys.add_driver(
+            FuelingRocket("fr", owner=self.sys, init=init, stop=stop, includes=includes, dt=dt)
+        )
+
+        self.sys.run_drivers()
+
+        data = self.sys.drivers["fr"].data
+        acel = np.asarray(data["rocket.a"])
+
+        np.testing.assert_allclose(acel[-2], 0.0, atol=10 ** (-10))
+        np.testing.assert_allclose(self.sys.rocket.stage_1.tank.weight_prop, 5.0, atol=10 ** (-10))
+        np.testing.assert_allclose(self.sys.g_tank.weight_prop, 5.0, atol=10 ** (-10))
+        np.testing.assert_allclose(self.sys.rocket.controller.flying, 1., atol=0.1)
+
+    def test_flight(self):
+
+        self.sys.drivers.clear()
+
+        dt = 0.1
+
+        init = {
+            "rocket.stage_1.tank.fuel.w_out_max": 3.0,
+            "g_tank.w_in": 0.0,
+            "g_tank.fuel.weight_p": 0.0,
+        }
+
+        stop = "rocket.stage_1.tank.weight_prop <= 0."
+
+        includes = ["rocket.a", "g_tank.weight", "rocket.stage_1.tank.weight_prop"]
+
+        self.sys.add_driver(
+            VerticalFlyingRocket("vfr", owner=self.sys, init=init, stop=stop, includes=includes, dt=dt)
+        )
+
+        self.sys.run_drivers()
+
+        data = self.sys.drivers["vfr"].data
+        data = data.drop(["Section", "Status", "Error code"], axis=1)
+
+        acel = np.asarray(data["rocket.a"])
+
+        np.testing.assert_allclose(acel[-2], 65.0, atol=10 ** (-10))
+        np.testing.assert_allclose(self.sys.rocket.stage_1.tank.weight_prop, 0.0, atol=10 ** (-10))
+        np.testing.assert_allclose(self.sys.g_tank.weight_prop, 0.0, atol=10 ** (-10))
diff --git a/rocket_twin/tests/test_flying_rocket.py b/rocket_twin/tests/test_flying_rocket.py
deleted file mode 100644
index 440603c..0000000
--- a/rocket_twin/tests/test_flying_rocket.py
+++ /dev/null
@@ -1,42 +0,0 @@
-import numpy as np
-
-from rocket_twin.drivers.vertical_flying_rocket import VerticalFlyingRocket
-from rocket_twin.systems import Station
-
-
-class TestVerticalFlyingRocket:
-    """Tests for the VerticalFlyingRocket driver."""
-
-    def test_run_once(self):
-        sys = Station("sys")
-        sys.rocket.stage_1.tank.fuel.weight_p = 5.0
-        dt = 0.1
-
-        init = {
-            "rocket.flying": True,
-            "controller.w_temp": 0.0,
-            "rocket.stage_1.controller.w_temp": 1.0,
-            "rocket.stage_1.tank.fuel.weight_p": 5.0,
-            "rocket.stage_1.tank.fuel.w_out_max": 3.0,
-            "g_tank.w_in": 0.0,
-            "g_tank.fuel.weight_p": 0.0,
-        }
-
-        stop = "rocket.stage_1.tank.weight_prop <= 0."
-
-        includes = ["rocket.a", "g_tank.weight", "rocket.stage_1.tank.weight_prop"]
-
-        sys.add_driver(
-            VerticalFlyingRocket("vfr", owner=sys, init=init, stop=stop, includes=includes, dt=dt)
-        )
-
-        sys.run_drivers()
-
-        data = sys.drivers["vfr"].data
-        data = data.drop(["Section", "Status", "Error code"], axis=1)
-
-        acel = np.asarray(data["rocket.a"])
-
-        np.testing.assert_allclose(acel[-2], 65.0, atol=10 ** (-10))
-        np.testing.assert_allclose(sys.rocket.stage_1.tank.weight_prop, 0.0, atol=10 ** (-10))
-        np.testing.assert_allclose(sys.g_tank.weight_prop, 0.0, atol=10 ** (-10))
diff --git a/rocket_twin/tests/test_mission.py b/rocket_twin/tests/test_mission.py
index 82b8eda..91fd857 100644
--- a/rocket_twin/tests/test_mission.py
+++ b/rocket_twin/tests/test_mission.py
@@ -13,8 +13,6 @@ def test_run_once(self):
 
         init = {
             "rocket.stage_1.tank.fuel.weight_p": 0.0,
-            "rocket.stage_1.controller.w_temp": 0.0,
-            "controller.w_temp": 1.0,
             "g_tank.fuel.weight_p": 10.0,
             "g_tank.w_in": 0.0,
             "g_tank.fuel.w_out_max": 3.0,
diff --git a/rocket_twin/tests/test_refuel_rocket.py b/rocket_twin/tests/test_refuel_rocket.py
deleted file mode 100644
index fcd5b02..0000000
--- a/rocket_twin/tests/test_refuel_rocket.py
+++ /dev/null
@@ -1,35 +0,0 @@
-import numpy as np
-
-from rocket_twin.drivers.fueling_rocket import FuelingRocket
-from rocket_twin.systems import Station
-
-
-class TestFuelingRocket:
-    """Tests for the FuelingRocket driver."""
-
-    def test_run_once(self):
-        sys = Station("sys")
-        dt = 0.1
-
-        init = {
-            "rocket.stage_1.tank.fuel.weight_p": 0.0,
-            "rocket.stage_1.controller.w_temp": 0.0,
-            "controller.w_temp": 1.0,
-            "g_tank.fuel.weight_p": 10.0,
-            "g_tank.w_in": 0.0,
-            "g_tank.fuel.w_out_max": 3.0,
-        }
-
-        stop = "rocket.stage_1.tank.weight_prop >= rocket.stage_1.tank.weight_max"
-
-        includes = ["rocket.a", "g_tank.weight", "rocket.stage_1.tank.weight_prop"]
-
-        sys.add_driver(
-            FuelingRocket("fr", owner=sys, init=init, stop=stop, includes=includes, dt=dt)
-        )
-
-        sys.run_drivers()
-
-        np.testing.assert_allclose(sys.rocket.a, 0.0, atol=10 ** (-10))
-        np.testing.assert_allclose(sys.rocket.stage_1.tank.weight_prop, 5.0, atol=10 ** (-10))
-        np.testing.assert_allclose(sys.g_tank.weight_prop, 5.0, atol=10 ** (-10))
diff --git a/rocket_twin/tests/test_stage.py b/rocket_twin/tests/test_stage.py
index dc908db..2854484 100644
--- a/rocket_twin/tests/test_stage.py
+++ b/rocket_twin/tests/test_stage.py
@@ -1,5 +1,5 @@
 import numpy as np
-from cosapp.drivers import RungeKutta, NonLinearSolver
+from cosapp.drivers import NonLinearSolver, RungeKutta
 from cosapp.recorders import DataFrameRecorder
 
 from rocket_twin.systems import Rocket, Station
@@ -29,7 +29,7 @@ def test_fuel(self):
         ]
 
         driver = self.sys.add_driver(RungeKutta("rk", order=4, dt=1))
-        solver = driver.add_child(NonLinearSolver('solver'))
+        solver = driver.add_child(NonLinearSolver("solver"))
         driver.time_interval = (0, 20)
         driver.set_scenario(init=init)
         driver.add_recorder(DataFrameRecorder(includes=includes), period=1.0)
@@ -37,8 +37,14 @@ def test_fuel(self):
         self.sys.run_drivers()
 
         data = driver.recorder.export_data()
-        data1 = data.drop(["Section", "Status", "Error code", "rocket.weight_prop_2", "rocket.weight_prop_3"], axis=1)
-        data2 = data.drop(["Section", "Status", "Error code", "g_tank.weight_prop", "rocket.weight_prop_1"], axis=1)
+        data1 = data.drop(
+            ["Section", "Status", "Error code", "rocket.weight_prop_2", "rocket.weight_prop_3"],
+            axis=1,
+        )
+        data2 = data.drop(
+            ["Section", "Status", "Error code", "g_tank.weight_prop", "rocket.weight_prop_1"],
+            axis=1,
+        )
 
         print(data1)
         print(data2)
@@ -66,7 +72,7 @@ def test_flight(self):
         ]
 
         driver = self.sys.add_driver(RungeKutta("rk", order=4, dt=1))
-        solver = driver.add_child(NonLinearSolver('solver'))
+        solver = driver.add_child(NonLinearSolver("solver"))
         driver.time_interval = (20, 40)
         driver.set_scenario(init=init, stop=stop)
         driver.add_recorder(DataFrameRecorder(includes=includes), period=1.0)
@@ -74,17 +80,23 @@ def test_flight(self):
         self.sys.run_drivers()
 
         data = driver.recorder.export_data()
-        data1 = data.drop(["Section", "Status", "Error code", "rocket.weight_prop_2", "rocket.weight_prop_3"], axis=1)
-        data2 = data.drop(["Section", "Status", "Error code", "g_tank.weight_prop", "rocket.weight_prop_1"], axis=1)
+        data1 = data.drop(
+            ["Section", "Status", "Error code", "rocket.weight_prop_2", "rocket.weight_prop_3"],
+            axis=1,
+        )
+        data2 = data.drop(
+            ["Section", "Status", "Error code", "g_tank.weight_prop", "rocket.weight_prop_1"],
+            axis=1,
+        )
 
         print(data1)
         print(data2)
 
-        #acel = np.asarray(data["rocket.a"])
-        #print(data)
+        # acel = np.asarray(data["rocket.a"])
+        # print(data)
 
         np.testing.assert_allclose(self.sys.rocket.weight_prop_1, 0.0, rtol=10 ** (-1))
         np.testing.assert_allclose(self.sys.rocket.weight_prop_2, 0.0, rtol=10 ** (-1))
         np.testing.assert_allclose(self.sys.rocket.weight_prop_3, 0.0, rtol=10 ** (-1))
         np.testing.assert_allclose(self.sys.rocket.geom.weight, 4.0, rtol=10 ** (-1))
-        #np.testing.assert_allclose(acel[-2], 40.0, atol=10 ** (-2))
+        # np.testing.assert_allclose(acel[-2], 40.0, atol=10 ** (-2))