Move exit code conversion to main function

2 files changed, 13 insertions, 15 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 4dea3cb..f3b2be4 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -7,7 +7,7 @@ import textwrap
 from pathlib import Path
 from subprocess import CompletedProcess
 from tempfile import NamedTemporaryFile
-from typing import Iterable, Tuple
+from typing import Iterable
 
 from snekbox import DEBUG
 
@@ -106,7 +106,7 @@ class NsJail:
                 log.error(msg)
 
     @staticmethod
-    def _consume_stdout(nsjail: subprocess.Popen) -> Tuple[int, str]:
+    def _consume_stdout(nsjail: subprocess.Popen) -> str:
         """
         Consume STDOUT, stopping when the output limit is reached or NsJail has exited.
 
@@ -116,7 +116,7 @@ class NsJail:
         is asked to terminate with a SIGKILL.
 
         Once the subprocess has exited, either naturally or because it was terminated,
-        we return the exit code as an int and the output as a single string.
+        we rturn the output as a single string.
         """
         output_size = 0
         output = []
@@ -129,19 +129,14 @@ class NsJail:
 
             if output_size > OUTPUT_MAX:
                 # Terminate the NsJail subprocess with SIGKILL.
-                log.info(f"Output exceeded the output limit, sending SIGKILL to NsJail.")
+                log.info("Output exceeded the output limit, sending SIGKILL to NsJail.")
                 nsjail.kill()
                 break
 
         # Ensure that we wait for the NsJail subprocess to terminate.
         nsjail.wait()
 
-        # When you send signal `N` to a subprocess to terminate it using Popen, it
-        # will return `-N` as its exit code. As we normally get `N + 128` back, we
-        # convert negative exit codes to the `N + 128` form.
-        returncode = -nsjail.returncode + 128 if nsjail.returncode < 0 else nsjail.returncode
-
-        return returncode, "".join(output)
+        return "".join(output)
 
     def python3(self, code: str) -> CompletedProcess:
         """Execute Python 3 code in an isolated environment and return the completed process."""
@@ -175,7 +170,12 @@ class NsJail:
             except ValueError:
                 return CompletedProcess(args, None, "ValueError: embedded null byte", None)
 
-            returncode, output = self._consume_stdout(nsjail)
+            output = self._consume_stdout(nsjail)
+
+            # When you send signal `N` to a subprocess to terminate it using Popen, it
+            # will return `-N` as its exit code. As we normally get `N + 128` back, we
+            # convert negative exit codes to the `N + 128` form.
+            returncode = -nsjail.returncode + 128 if nsjail.returncode < 0 else nsjail.returncode
 
             log_lines = nsj_log.read().decode("utf-8").splitlines()
             if not log_lines and returncode == 255:
diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py
index f4e678c..40c27f9 100644
--- a/tests/test_nsjail.py
+++ b/tests/test_nsjail.py
@@ -5,7 +5,7 @@ import unittest
 import unittest.mock
 from textwrap import dedent
 
-from snekbox.nsjail import MEM_MAX, NsJail, READ_CHUNK_SIZE, OUTPUT_MAX
+from snekbox.nsjail import MEM_MAX, NsJail, OUTPUT_MAX, READ_CHUNK_SIZE
 
 
 class NsJailTests(unittest.TestCase):
@@ -195,9 +195,7 @@ class NsJailTests(unittest.TestCase):
 
         # Go 10 chunks over to make sure we exceed the limit
         nsjail_subprocess.stdout = io.StringIO((expected_chunks + 10) * chunk)
-        nsjail_subprocess.returncode = -9
         nsjail_subprocess.poll.return_value = None
 
-        returncode, output = self.nsjail._consume_stdout(nsjail_subprocess)
-        self.assertEqual(returncode, 137)
+        output = self.nsjail._consume_stdout(nsjail_subprocess)
         self.assertEqual(output, chunk * expected_chunks)