diff --git a/reflex/testing.py b/reflex/testing.py
index 6c2e4d884..9b9ea1c06 100644
--- a/reflex/testing.py
+++ b/reflex/testing.py
@@ -189,7 +189,9 @@ class AppHarness:
         if isinstance(app_source, str):
             return app_source
         source = inspect.getsource(app_source)
-        source = re.sub(r"^\s*def\s+\w+\s*\(.*?\):", "", source, flags=re.DOTALL)
+        source = re.sub(
+            r"^\s*def\s+\w+\s*\(.*?\)(\s+->\s+\w+)?:", "", source, flags=re.DOTALL
+        )
         return textwrap.dedent(source)
 
     def _initialize_app(self):