Handle BaseException In Code Question

Change the exception handler in the unittest template to handle
BaseExceptions that don't inherit from Exception. Additionally, account
for user invoked system exits.

Signed-off-by: Hassan Abouelela <[email protected]>

1 files changed, 11 insertions, 4 deletions

diff --git a/resources/unittest_template.py b/resources/unittest_template.py
index 990d077..05730ce 100644
--- a/resources/unittest_template.py
+++ b/resources/unittest_template.py
@@ -19,6 +19,9 @@ class RunnerTestCase(unittest.TestCase):
 ### UNIT CODE
 
 
+normal_exit = False
+
+
 def _exit_sandbox(code: int) -> NoReturn:
     """
     Exit the sandbox by printing the result to the actual stdout and exit with the provided code.
@@ -32,6 +35,8 @@ def _exit_sandbox(code: int) -> NoReturn:
     137 can also be generated by NsJail when killing the process.
     """
     print(RESULT.getvalue(), file=ORIGINAL_STDOUT, end="")
+    global normal_exit
+    normal_exit = True
     sys.exit(code)
 
 
@@ -81,10 +86,12 @@ try:
     # Load the user code as a global module variable
     try:
         module = _load_user_module()
-    except Exception:
-        RESULT.write("Uncaught exception while loading user code.")
+    except BaseException as e:
+        RESULT.write(f"Uncaught exception while loading user code: {e}")
         _exit_sandbox(6)
     _main()
-except Exception:
-    RESULT.write("Uncaught exception inside runner.")
+except BaseException as e:
+    if isinstance(e, SystemExit) and normal_exit:
+        raise e from None
+    RESULT.write(f"Uncaught exception inside runner: {e}")
     _exit_sandbox(99)