From 6ecc16fc4bec36698623091c50c905b88de6dcc9 Mon Sep 17 00:00:00 2001
From: Simon Tatham <anakin@pobox.com>
Date: Sun, 24 Mar 2019 09:56:57 +0000
Subject: [PATCH] cryptsuite: clean up exit handling.

Now we only run the final memory-leak check if we didn't already have
some other error to report, or some other exception that terminated
the process.

Also, we wait for the subprocess to terminate before returning control
to the shell, so that any last-minute complaints from Leak Sanitiser
appear before rather than after the shell prompt comes back.

While I'm here, I've also made check_return_status tolerate the case
in which the child process never got started at all. That way, if a
failure manages to occur before even getting _that_ far, there won't
be a cascade failure from check_return_status getting confused
afterwards.
---
 test/cryptsuite.py | 21 ++++++++++++++-------
 test/testcrypt.py  | 15 ++++++++++-----
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/test/cryptsuite.py b/test/cryptsuite.py
index 773526fd..b3dbee3d 100755
--- a/test/cryptsuite.py
+++ b/test/cryptsuite.py
@@ -1833,10 +1833,17 @@ class standard_test_vectors(MyTestBase):
             self.assertEqual(crc32_rfc1662(vec), 0x2144DF1C)
 
 if __name__ == "__main__":
-    try:
-        unittest.main()
-    finally:
-        # On exit, make sure we check the subprocess's return status,
-        # so that if Leak Sanitiser detected any memory leaks, the
-        # test will turn into a failure at the last minute.
-        childprocess.check_return_status()
+    # Run the tests, suppressing automatic sys.exit and collecting the
+    # unittest.TestProgram instance returned by unittest.main instead.
+    testprogram = unittest.main(exit=False)
+
+    # If any test failed, just exit with failure status.
+    if not testprogram.result.wasSuccessful():
+        childprocess.wait_for_exit()
+        sys.exit(1)
+
+    # But if no tests failed, we have one last check to do: look at
+    # the subprocess's return status, so that if Leak Sanitiser
+    # detected any memory leaks, the success return status will turn
+    # into a failure at the last minute.
+    childprocess.check_return_status()
diff --git a/test/testcrypt.py b/test/testcrypt.py
index 95d311d9..659ea4ca 100644
--- a/test/testcrypt.py
+++ b/test/testcrypt.py
@@ -26,6 +26,7 @@ class ChildProcess(object):
     def __init__(self):
         self.sp = None
         self.debug = None
+        self.exitstatus = None
 
         dbg = os.environ.get("PUTTY_TESTCRYPT_DEBUG")
         if dbg is not None:
@@ -62,12 +63,16 @@ class ChildProcess(object):
             unicode_to_bytes(arg) for arg in args))
         argcount = int(self.read_line())
         return [self.read_line() for arg in range(argcount)]
+    def wait_for_exit(self):
+        if self.sp is not None:
+            self.sp.stdin.close()
+            self.exitstatus = self.sp.wait()
+            self.sp = None
     def check_return_status(self):
-        assert self.sp is not None
-        self.sp.stdin.close()
-        status = self.sp.wait()
-        if status != 0:
-            raise Exception("testcrypt returned exit status {}".format(status))
+        self.wait_for_exit()
+        if self.exitstatus is not None and self.exitstatus != 0:
+            raise Exception("testcrypt returned exit status {}"
+                            .format(self.exitstatus))
 
 childprocess = ChildProcess()