Cleanup sys.exit, os._exit, flush stdout, fix tests (#2931)

* add flush to sys.exit patch + remove flush in unneeded places

* add test for required args

* lint

* remove os._exit in event detection

* fix failed tests

* fix failed tests

* argparse no exit?

* fix argparse exit

* fix argparse exit

* lint

* lint

* fix test detaching

gdbinit.py

@@ -169,7 +169,6 @@ def check_doubleload():
         print(
             "To fix this, please remove the line 'source your-path/gdbinit.py' from your .gdbinit file."
         )
-        sys.stdout.flush()
         sys.exit(1)
 
 
@@ -180,7 +179,19 @@ def rewire_exit():
         # a segfault. See:
         # https://github.com/pwndbg/pwndbg/pull/2900#issuecomment-2825456636
         # https://sourceware.org/bugzilla/show_bug.cgi?id=31946
-        sys.exit = os._exit
+        def _patched_exit(exit_code):
+            # argparse requires a SystemExit exception, otherwise our CLI commands will exit incorrectly on invalid arguments
+            stack_list = traceback.extract_stack(limit=2)
+            if len(stack_list) == 2:
+                p = stack_list[0]
+                if p.filename.endswith("/argparse.py"):
+                    raise SystemExit()
+
+            sys.stdout.flush()
+            sys.stderr.flush()
+            os._exit(exit_code)
+
+        sys.exit = _patched_exit
 
 
 def main() -> None:
@@ -201,7 +212,6 @@ def main() -> None:
         venv_path = get_venv_path(src_root)
         if not venv_path.exists():
             print(f"Cannot find Pwndbg virtualenv directory: {venv_path}. Please re-run setup.sh")
-            sys.stdout.flush()
             sys.exit(1)
         no_auto_update = os.getenv("PWNDBG_NO_AUTOUPDATE")
         if no_auto_update is None:
@@ -248,5 +258,4 @@ try:
 
 except Exception:
     print(traceback.format_exc(), file=sys.stderr)
-    sys.stdout.flush()
     sys.exit(1)

pwndbg/gdblib/events.py

@@ -169,9 +169,7 @@ To address this, you have three options:
 """
         )
     )
-    import os
-
-    os._exit(1)
+    sys.exit(1)
 
 
 def wrap_safe_event_handler(event_handler: Callable[P, T], event_type: Any) -> Callable[P, T]:

tests/pytests_collect.py

@@ -9,7 +9,6 @@ TESTS_PATH = os.environ.get("TESTS_PATH")
 
 if TESTS_PATH is None:
     print("'TESTS_PATH' environment variable not set. Failed to collect tests.")
-    sys.stdout.flush()
     sys.exit(1)
 
 
@@ -29,7 +28,6 @@ rv = pytest.main(["--collect-only", TESTS_PATH], plugins=[collector])
 
 if rv == pytest.ExitCode.INTERRUPTED:
     print("Failed to collect all tests, perhaps there is a syntax error in one of test files?")
-    sys.stdout.flush()
     sys.exit(1)
 
 
@@ -38,5 +36,4 @@ for nodeid in collector.collected:
     print("Test:", nodeid)
 
 # easy way to exit GDB session
-sys.stdout.flush()
 sys.exit(0)

tests/pytests_launcher.py

@@ -37,5 +37,9 @@ if (cov := coverage.Coverage.current()) is not None:
     cov.stop()
     cov.save()
 
+# `sys.exit` triggers a GDB detach, while `os._exit` does not.
+# This allows the debugging session to remain at the same PC location,
+# which is useful for attaching to qemu-system multiple times.
 sys.stdout.flush()
-sys.exit(return_code)
+sys.stderr.flush()
+os._exit(return_code)

tests/qemu-tests/conftest.py

@@ -32,7 +32,6 @@ def qemu_assembly_run():
 
     if QEMU_PORT is None:
         print("'QEMU_PORT' environment variable not set")
-        sys.stdout.flush()
         sys.exit(1)
 
     def _start_binary(asm: str, arch: str, endian: Literal["big", "little"] | None = None):
@@ -103,7 +102,6 @@ def qemu_start_binary():
 
     if QEMU_PORT is None:
         print("'QEMU_PORT' environment variable not set")
-        sys.stdout.flush()
         sys.exit(1)
 
     def _start_binary(path: str, arch: str, endian: Literal["big", "little"] | None = None):

tests/tests.py

@@ -173,29 +173,34 @@ class TestStats:
 
     def handle_test_result(self, test_result: TEST_RETURN_TYPE, args, test_dir_path):
         (process, test_case, duration) = test_result
-        content = process.stdout
-
-        # Extract the test name and result using regex
-        testname = re.search(rf"^({test_dir_path}/[^ ]+)", content, re.MULTILINE)[0]
-        result = re.search(
-            r"(\x1b\[3.m(PASSED|FAILED|SKIPPED|XPASS|XFAIL)\x1b\[0m)", content, re.MULTILINE
-        )[0]
-
-        (_, testname) = testname.split("::")
-
-        if "FAIL" in result:
+        if args.serial:
+            # Serial mode does not capture stdout, so it's not possible to check the result
+            return
+
+        test_status = "FAIL"
+        if process.returncode == 0:
+            result = re.search(
+                r"(\x1b\[3.m(PASSED|FAILED|SKIPPED|XPASS|XFAIL)\x1b\[0m)",
+                process.stdout,
+                re.MULTILINE,
+            )
+            if result:
+                test_status = result[0]
+
+        if "FAIL" in test_status:
             self.fail_tests += 1
             self.fail_tests_names.append(test_case)
-        elif "PASS" in result:
+        elif "PASS" in test_status:
             self.pass_tests += 1
-        elif "SKIP" in result:
+        elif "SKIP" in test_status:
             self.skip_tests += 1
-        print(f"{testname:<70} {result} {duration:.2f}s")
+        print(f"{test_case:<70} {test_status} {duration:.2f}s")
 
         # Only show the output of failed tests unless the verbose flag was used
-        if args.verbose or "FAIL" in result:
+        if args.verbose or "FAIL" in test_status:
             print("")
-            print(content)
+            print(process.stderr)
+            print(process.stdout)
 
 
 def run_tests_and_print_stats(