diff --git a/CADETPythonSimulator/unit_operation.py b/CADETPythonSimulator/unit_operation.py
index 521b200..3da1e35 100644
--- a/CADETPythonSimulator/unit_operation.py
+++ b/CADETPythonSimulator/unit_operation.py
@@ -327,7 +327,7 @@ def get_outlet_state(
     def get_outlet_state_flat(
             self,
             unit_port_index: int
-            ) -> NoReturn:
+            ) -> dict[str, np.ndarray]:
         """
         Get the state of the unit operation outlet for a given port.
 
diff --git a/tests/test_rejection.py b/tests/test_rejection.py
index d1c1aac..6ec773d 100644
--- a/tests/test_rejection.py
+++ b/tests/test_rejection.py
@@ -6,12 +6,6 @@
 )
 
 
-
-
-
-
-
-
 TestCaseCutof = {
             "model": StepCutOff,
             "model_param": {
@@ -30,7 +24,9 @@
 
 class TestRejection():
     def test_get_rejection(self, parameters):
-
+        """
+        Test to check wheter the get_rejection function works as intended
+        """
         model = parameters["model"](**parameters["model_param"])
 
         solution = [model.get_rejection(weight) for weight in parameters["weights"]]
diff --git a/tests/test_unit_operation.py b/tests/test_unit_operation.py
index b4698fd..115c18e 100644
--- a/tests/test_unit_operation.py
+++ b/tests/test_unit_operation.py
@@ -189,7 +189,7 @@ def __init__(self,
                 'n_outlet_ports': 1,
                 'n_dof': 10,
                 'states': {
-                    'retentate': [0., 1., 2., 3., 4., 5.],
+                    'cake': [0., 1., 2., 3., 4., 5.],
                     'permeate': [6., 7., 8., 9.],
                 },
                 'inlet_state': {