diff --git a/monkey/infection_monkey/utils/capture_output.py b/monkey/infection_monkey/utils/capture_output.py
index 898bd6a7e..f509fa908 100644
--- a/monkey/infection_monkey/utils/capture_output.py
+++ b/monkey/infection_monkey/utils/capture_output.py
@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import io
 import sys
 
 
 class StdoutCapture:
-    def __enter__(self) -> None:
+    def __enter__(self) -> StdoutCapture:
         self._orig_stdout = sys.stdout
         self._new_stdout = io.StringIO()
         sys.stdout = self._new_stdout