diff --git a/paseos/tests/activity_test.py b/paseos/tests/activity_test.py
index 969b07b2..8e135892 100644
--- a/paseos/tests/activity_test.py
+++ b/paseos/tests/activity_test.py
@@ -1,5 +1,5 @@
 """Simple test of starting an activity"""
-from test_utils import get_default_instance, wait_for_activity
+from test_utils import get_default_instance
 
 from paseos import SpacecraftActor
 import asyncio
@@ -58,7 +58,7 @@ async def func(args):
 
     # Run the activity
     sim.perform_activity("Testing", activity_func_args=[test_val])
-    await wait_for_activity(sim)
+    await sim.wait_for_activity()
 
     # Check activity result
     assert test_val[0] == 10
@@ -99,7 +99,7 @@ async def act2(args):
 
     # try running it
     sim.perform_activity("act1", activity_func_args=[test_value])
-    await wait_for_activity(sim)
+    await sim.wait_for_activity()
 
     # Value should be 42 as first activity is started but then an error occurs trying the second
     assert test_value[0] == 42
@@ -150,7 +150,7 @@ async def on_termination(args):
         constraint_func_args=[sat1],
         termination_func_args=[test_value, test_value2],
     )
-    await wait_for_activity(sim)
+    await sim.wait_for_activity()
 
     assert test_value == test_value2
     assert sat1.battery_level_in_Ws >= 505
diff --git a/paseos/tests/operations_monitor_test.py b/paseos/tests/operations_monitor_test.py
index 0622f5fc..615e5c90 100644
--- a/paseos/tests/operations_monitor_test.py
+++ b/paseos/tests/operations_monitor_test.py
@@ -4,7 +4,6 @@
 import pytest
 import pykep as pk
 
-from test_utils import wait_for_activity
 import paseos
 from paseos import ActorBuilder, SpacecraftActor, load_default_cfg
 
@@ -43,10 +42,10 @@ async def func2(args):
 
     # Run the activity
     sim.perform_activity("Activity_1")
-    await wait_for_activity(sim)
+    await sim.wait_for_activity()
 
     sim.perform_activity("Activity_2")
-    await wait_for_activity(sim)
+    await sim.wait_for_activity()
 
     # Try out item function
     sim.monitor["state_of_charge"]
diff --git a/paseos/tests/radiation_test.py b/paseos/tests/radiation_test.py
index b5fd5993..1012baa0 100644
--- a/paseos/tests/radiation_test.py
+++ b/paseos/tests/radiation_test.py
@@ -4,7 +4,7 @@
 import numpy as np
 import pytest
 
-from test_utils import get_default_instance, wait_for_activity
+from test_utils import get_default_instance
 from paseos.radiation.radiation_model import RadiationModel
 from paseos import ActorBuilder
 import paseos
@@ -83,7 +83,7 @@ async def func(args):
 
     # Run the activity
     sim.perform_activity("Testing")
-    await wait_for_activity(sim)
+    await sim.wait_for_activity()
 
     # Check sat1 died.
     assert sat1.is_dead
@@ -120,7 +120,7 @@ async def func(args):
 
     # Run the activity
     sim.perform_activity("Testing", activity_func_args=[test_val])
-    await wait_for_activity(sim)
+    await sim.wait_for_activity()
 
     # Check activity result, exact result is runtime dependent.
     assert test_val[0] > 4 and test_val[0] < 8
diff --git a/paseos/tests/thermal_model_test.py b/paseos/tests/thermal_model_test.py
index 0744a6e3..46084044 100644
--- a/paseos/tests/thermal_model_test.py
+++ b/paseos/tests/thermal_model_test.py
@@ -1,7 +1,6 @@
 """Simple test of the thermal model to see if temperatures evolve as expected"""
 import pykep as pk
 
-from test_utils import wait_for_activity
 import paseos
 from paseos import SpacecraftActor, ActorBuilder, load_default_cfg
 import asyncio
@@ -51,7 +50,7 @@ async def func(args):
 
     # Run the activity
     sim.perform_activity("Activity_1")
-    await wait_for_activity(sim)
+    await sim.wait_for_activity()
     assert sat1.temperature_in_K > 285
     assert sat1.temperature_in_K < 300
     sim.save_status_log_csv("thermal_test.csv")
diff --git a/paseos/tests/time_multiplier_test.py b/paseos/tests/time_multiplier_test.py
index 1623a5ab..bb04ce42 100644
--- a/paseos/tests/time_multiplier_test.py
+++ b/paseos/tests/time_multiplier_test.py
@@ -4,7 +4,6 @@
 import pytest
 import pykep as pk
 
-from test_utils import wait_for_activity
 import paseos
 from paseos import ActorBuilder, SpacecraftActor, load_default_cfg
 
@@ -43,7 +42,7 @@ async def func(args):
 
     # Run the activity
     sim.perform_activity("Testing", activity_func_args=[test_val])
-    await wait_for_activity(sim)
+    await sim.wait_for_activity()
 
     # Check activity result
     assert test_val[0] == 5