🐛 Fix attach debugger when chain not connected

wake/cli/test.py

@@ -95,7 +95,12 @@ def run_no_pytest(
 
         try:
             for _, func in test_functions:
-                func()
+                try:
+                    func()
+                except Exception:
+                    if debug:
+                        attach_debugger(*sys.exc_info())
+                    raise
                 reset_exception_handled()
         finally:
             if coverage:

wake/development/core.py

@@ -1590,7 +1590,7 @@ def _connect(
             if not isinstance(e, BdbQuit):
                 exception_handler = get_exception_handler()
                 if exception_handler is not None:
-                    exception_handler(e)
+                    exception_handler(*sys.exc_info())
                 raise
         finally:
             self._connect_finalize()
@@ -1752,7 +1752,7 @@ def change_automine(self, automine: bool):
             if not isinstance(e, BdbQuit):
                 exception_handler = get_exception_handler()
                 if exception_handler is not None:
-                    exception_handler(e)
+                    exception_handler(*sys.exc_info())
                 raise
         finally:
             self._chain_interface.set_automine(automine_was)
@@ -1835,7 +1835,7 @@ def snapshot_and_revert(self):
             if not isinstance(e, BdbQuit):
                 exception_handler = get_exception_handler()
                 if exception_handler is not None:
-                    exception_handler(e)
+                    exception_handler(*sys.exc_info())
                 raise
         finally:
             self.revert(snapshot_id)

wake/development/globals.py

@@ -4,9 +4,19 @@
 from collections import defaultdict
 from pathlib import Path
 from types import TracebackType
-from typing import TYPE_CHECKING, Callable, DefaultDict, List, Optional, Set, Tuple
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    DefaultDict,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+)
 from urllib.error import HTTPError
 
+import rich.traceback
 import rich_click
 from ipdb.__main__ import _init_pdb
 
@@ -24,7 +34,16 @@
 
 
 # must be declared before functions that use it because of a bug in Python (https://bugs.python.org/issue34939)
-_exception_handler: Optional[Callable[[Exception], None]] = None
+_exception_handler: Optional[
+    Callable[
+        [
+            Optional[Type[BaseException]],
+            Optional[BaseException],
+            Optional[TracebackType],
+        ],
+        None,
+    ]
+] = None
 _exception_handled = False
 
 _coverage_handler: Optional[CoverageHandler] = None
@@ -33,21 +52,27 @@
 _verbosity: int = 0
 
 
-def attach_debugger(e: Exception):
+def attach_debugger(
+    e_type: Optional[Type[BaseException]],
+    e: Optional[BaseException],
+    tb: Optional[TracebackType],
+):
     global _exception_handled
 
     if _exception_handled:
         return
     _exception_handled = True
 
-    import sys
     import traceback
 
     from wake.cli.console import console
 
-    tb: Optional[TracebackType] = sys.exc_info()[2]
+    assert e_type is not None
+    assert e is not None
     assert tb is not None
-    console.print_exception()
+
+    rich_tb = rich.traceback.Traceback.from_exception(e_type, e, tb)
+    console.print(rich_tb)
 
     frames = []
 
@@ -71,11 +96,29 @@ def attach_debugger(e: Exception):
     p.interaction(None, tb)
 
 
-def get_exception_handler() -> Optional[Callable[[Exception], None]]:
+def get_exception_handler() -> Optional[
+    Callable[
+        [
+            Optional[Type[BaseException]],
+            Optional[BaseException],
+            Optional[TracebackType],
+        ],
+        None,
+    ]
+]:
     return _exception_handler
 
 
-def set_exception_handler(handler: Callable[[Exception], None]):
+def set_exception_handler(
+    handler: Callable[
+        [
+            Optional[Type[BaseException]],
+            Optional[BaseException],
+            Optional[TracebackType],
+        ],
+        None,
+    ]
+):
     global _exception_handler
     _exception_handler = handler
 

wake/testing/fuzzing/fuzzer.py

@@ -10,7 +10,7 @@
 import types
 from contextlib import redirect_stderr, redirect_stdout
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterable, List, Optional
+from typing import Any, Callable, Dict, Iterable, List, Optional, Type
 
 import rich.progress
 from pathvalidate import sanitize_filename  # type: ignore
@@ -35,26 +35,6 @@
 from wake.utils.tee import StderrTee, StdoutTee
 
 
-def _run_core(
-    fuzz_test: Callable,
-    index: int,
-    random_seed: bytes,
-    finished_event: multiprocessing.synchronize.Event,
-    err_child_conn: multiprocessing.connection.Connection,
-    cov_child_conn: multiprocessing.connection.Connection,
-    coverage: Optional[CoverageHandler],
-):
-    console.print(f"Using random seed '{random_seed.hex()}' for process #{index}")
-
-    fuzz_test()
-
-    err_child_conn.send(None)
-    if coverage is not None:
-        # final coverage update
-        cov_child_conn.send(coverage.get_contract_ide_coverage())
-    finished_event.set()
-
-
 def _run(
     fuzz_test: Callable,
     index: int,
@@ -66,29 +46,38 @@ def _run(
     cov_child_conn: multiprocessing.connection.Connection,
     coverage: Optional[CoverageHandler],
 ):
-    def exception_handler(e: Exception) -> None:
+    def exception_handler(
+        e_type: Optional[Type[BaseException]],
+        e: Optional[BaseException],
+        tb: Optional[types.TracebackType],
+    ) -> None:
         for ctx_manager in ctx_managers:
             ctx_manager.__exit__(None, None, None)
         ctx_managers.clear()
 
-        exc_info = sys.exc_info()
+        assert e_type is not None
+        assert e is not None
+        assert tb is not None
+
+        nonlocal exception_handled
+        exception_handled = True
+
         try:
-            pickled = pickle.dumps(exc_info)
+            pickled = pickle.dumps((e_type, e, tb))
         except Exception:
-            pickled = pickle.dumps(
-                (exc_info[0], Exception(repr(exc_info[1])), exc_info[2])
-            )
+            pickled = pickle.dumps((e_type, Exception(repr(e)), tb))
         err_child_conn.send(pickled)
         finished_event.set()
 
         try:
             attach: bool = err_child_conn.recv()
             if attach:
                 sys.stdin = os.fdopen(0)
-                attach_debugger(e)
+                attach_debugger(e_type, e, tb)
         finally:
             finished_event.set()
 
+    exception_handled = False
     last_coverage_sync = time.perf_counter()
 
     def coverage_callback() -> None:
@@ -122,15 +111,20 @@ def coverage_callback() -> None:
         for ctx_manager in ctx_managers:
             ctx_manager.__enter__()
 
-        _run_core(
-            fuzz_test,
-            index,
-            random_seed,
-            finished_event,
-            err_child_conn,
-            cov_child_conn,
-            coverage,
-        )
+        console.print(f"Using random seed '{random_seed.hex()}' for process #{index}")
+
+        try:
+            fuzz_test()
+        except Exception:
+            if not exception_handled:
+                exception_handler(*sys.exc_info())
+            raise
+
+        err_child_conn.send(None)
+        if coverage is not None:
+            # final coverage update
+            cov_child_conn.send(coverage.get_contract_ide_coverage())
+        finished_event.set()
     except Exception:
         pass
     finally:

wake/testing/pytest_plugin_multiprocess.py

@@ -9,7 +9,8 @@
 import time
 from contextlib import redirect_stderr, redirect_stdout
 from pathlib import Path
-from typing import List, Optional
+from types import TracebackType
+from typing import List, Optional, Type
 
 import pytest
 from pathvalidate import sanitize_filename
@@ -36,6 +37,7 @@ class PytestWakePluginMultiprocess:
     _random_seed: bytes
     _tee: bool
     _debug: bool
+    _exception_handled: bool
 
     _ctx_managers: List
 
@@ -58,6 +60,7 @@ def __init__(
         self._random_seed = random_seed
         self._tee = tee
         self._debug = debug
+        self._exception_handled = False
 
         self._ctx_managers = []
 
@@ -77,6 +80,34 @@ def _cleanup_stdio(self):
             ctx_manager.__exit__(None, None, None)
         self._ctx_managers.clear()
 
+    def _exception_handler(
+        self,
+        e_type: Optional[Type[BaseException]],
+        e: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
+        self._cleanup_stdio()
+        self._exception_handled = True
+
+        assert e_type is not None
+        assert e is not None
+        assert tb is not None
+
+        try:
+            pickled = pickle.dumps((e_type, e, tb))
+        except Exception:
+            pickled = pickle.dumps((e_type, Exception(repr(e)), tb))
+        self._queue.put(("exception", self._index, pickled), block=True)
+
+        attach: bool = self._conn.recv()
+        try:
+            if attach:
+                sys.stdin = os.fdopen(0)
+                attach_debugger(e_type, e, tb)
+        finally:
+            self._setup_stdio()
+            self._conn.send(("exception_handled",))
+
     def pytest_configure(self, config: pytest.Config):
         self._f = open(self._log_file, "w")
         self._setup_stdio()
@@ -95,6 +126,7 @@ def pytest_collection_finish(self, session: Session):
 
     def pytest_runtest_setup(self, item):
         reset_exception_handled()
+        self._exception_handled = False
 
     def pytest_internalerror(
         self, excrepr, excinfo: pytest.ExceptionInfo[BaseException]
@@ -107,6 +139,12 @@ def pytest_internalerror(
             )
         self._queue.put(("pytest_internalerror", self._index, pickled), block=True)
 
+    def pytest_exception_interact(self, node, call, report):
+        if self._debug and not self._exception_handled:
+            self._exception_handler(
+                call.excinfo.type, call.excinfo.value, call.excinfo.tb
+            )
+
     def pytest_runtestloop(self, session: Session):
         if (
             session.testsfailed
@@ -120,27 +158,6 @@ def pytest_runtestloop(self, session: Session):
         if session.config.option.collectonly:
             return True
 
-        def exception_handler(e: Exception) -> None:
-            self._cleanup_stdio()
-
-            exc_info = sys.exc_info()
-            try:
-                pickled = pickle.dumps(exc_info)
-            except Exception:
-                pickled = pickle.dumps(
-                    (exc_info[0], Exception(repr(exc_info[1])), exc_info[2])
-                )
-            self._queue.put(("exception", self._index, pickled), block=True)
-
-            attach: bool = self._conn.recv()
-            try:
-                if attach:
-                    sys.stdin = os.fdopen(0)
-                    attach_debugger(e)
-            finally:
-                self._setup_stdio()
-                self._conn.send(("exception_handled",))
-
         last_coverage_sync = time.perf_counter()
 
         def coverage_callback() -> None:
@@ -167,7 +184,7 @@ def signal_handler(sig, frame):
         signal.signal(signal.SIGTERM, signal_handler)
 
         if self._debug:
-            set_exception_handler(exception_handler)
+            set_exception_handler(self._exception_handler)
         if self._coverage is not None:
             set_coverage_handler(self._coverage)
             self._coverage.set_callback(coverage_callback)

wake/testing/pytest_plugin_single.py

@@ -41,6 +41,10 @@ def __init__(
     def pytest_runtest_setup(self, item):
         reset_exception_handled()
 
+    def pytest_exception_interact(self, node, call, report):
+        if self._debug:
+            attach_debugger(call.excinfo.type, call.excinfo.value, call.excinfo.tb)
+
     def pytest_runtestloop(self, session: Session):
         if (
             session.testsfailed