diff --git a/tests/unit_tests/test_axon.py b/tests/unit_tests/test_axon.py
index 5dac992d21..20fdee4369 100644
--- a/tests/unit_tests/test_axon.py
+++ b/tests/unit_tests/test_axon.py
@@ -16,6 +16,7 @@
 # DEALINGS IN THE SOFTWARE.
 
 
+import contextlib
 import re
 import threading
 import time
@@ -795,6 +796,7 @@ async def test_threaded_fastapi():
     server_started = threading.Event()
     server_stopped = threading.Event()
 
+    @contextlib.asynccontextmanager
     async def lifespan(app):
         server_started.set()
         yield