✨ add custom pdb and custom debug trigger

wake-coverage.cov

@@ -0,0 +1,4 @@
+{
+    "version": "1.0",
+    "data": {}
+}

wake/testing/custom_pdb.py

@@ -0,0 +1,37 @@
+import pdb
+import sys
+
+class CustomPdb(pdb.Pdb):
+    def __init__(self, prev_stdin, conn, *args, **kwargs):
+        """
+        Custom Pdb constructor to accept the stdin and connection.
+        """
+        super().__init__(*args, **kwargs)  # Initialize the base Pdb class
+        self.prev_stdin = prev_stdin  # Store the original stdin
+        self.conn = conn  # Store the connection for parent communication
+
+    def do_continue(self, arg):
+        """
+        Override the 'continue' (c) command to perform cleanup after continuing.
+        """
+        self.cleanup_before_exit()
+        return super().do_continue(arg)
+
+    do_c = do_cont = do_continue
+
+    def do_quit(self, arg):
+        """
+        Override the 'quit' (q) command to perform cleanup before quitting.
+        """
+        self.cleanup_before_exit()
+        return super().do_quit(arg)
+
+    do_q = do_exit = do_quit 
+
+    def cleanup_before_exit(self):
+        """
+        This function performs the cleanup before exiting the debugger.
+        """
+        print("Performing cleanup before exiting the debugger...")
+        sys.stdin = self.prev_stdin  # Restore stdin to its original state
+        self.conn.send(("breakpoint_handled",))  

wake/testing/pytest_plugin_multiprocess.py

@@ -25,9 +25,12 @@
     set_coverage_handler,
     set_exception_handler,
 )
+from ipdb.__main__ import _init_pdb
 from wake.testing.coverage import CoverageHandler
 from wake.utils.tee import StderrTee, StdoutTee
 
+from wake.testing.custom_pdb import CustomPdb
+
 
 class PytestWakePluginMultiprocess:
     _index: int
@@ -177,6 +180,22 @@ def coverage_callback() -> None:
                 except queue.Full:
                     pass
 
+        def custom_debugger():
+            self._cleanup_stdio()
+            self._queue.put(("breakpoint", self._index), block=True)
+            attach: bool = self._conn.recv()
+            if attach:
+                prev = sys.stdin
+                sys.stdin = os.fdopen(0)
+                frame = sys._getframe(1)
+                p = CustomPdb(prev, self._conn)
+                p.set_trace(frame)
+            else:
+                # trace nothing, same as continue
+                self._conn.send(("breakpoint_handled",))  
+
+        sys.breakpointhook = custom_debugger
+
         def signal_handler(sig, frame):
             raise KeyboardInterrupt()
 

wake/testing/pytest_plugin_multiprocess_server.py

@@ -244,6 +244,34 @@ def pytest_runtestloop(self, session: pytest.Session):
                         )
                         if progress is not None:
                             progress.start()
+                    elif msg[0] == "breakpoint":
+                        if progress is not None:
+                            progress.stop()
+                            console.print(
+                                f"Process #{index} reach breakpoint."
+                            )
+
+                            attach = None
+                            while attach is None:
+                                response = input(
+                                    "Would you like to attach the debugger? [y/n] ('n' to continue operations)"
+                                )
+                                if response == "y":
+                                    attach = True
+                                elif response == "n":
+                                    attach = False
+                        else:
+                            attach = False
+
+                        self._processes[index][1].send(attach)
+
+                        # wait for debugger to finish
+                        assert self._processes[index][1].recv() == (
+                            "breakpoint_handled",
+                        )
+                        if progress is not None:
+                            progress.start()
+
                     elif msg[0] == "pytest_runtest_logreport":
                         report: pytest.TestReport = msg[2]
                         reports.append(report)