Add test for stdout output truncation

I've added a test that checks if output exceeding the limit is
correctly truncated. To make the test more robust, I've defined a
constant for the read chunk size.

2 files changed, 21 insertions, 2 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 8edf546..4dea3cb 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -29,6 +29,7 @@ MEM_MAX = 52428800
 
 # Limit of stdout bytes we consume before terminating nsjail
 OUTPUT_MAX = 1_000_000  # 1 MB
+READ_CHUNK_SIZE = 10_000  # chars
 
 
 class NsJail:
@@ -122,7 +123,7 @@ class NsJail:
 
         # We'll consume STDOUT as long as the NsJail subprocess is running.
         while nsjail.poll() is None:
-            chars = nsjail.stdout.read(10_000)
+            chars = nsjail.stdout.read(READ_CHUNK_SIZE)
             output_size += sys.getsizeof(chars)
             output.append(chars)
 
diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py
index 691eb5b..f4e678c 100644
--- a/tests/test_nsjail.py
+++ b/tests/test_nsjail.py
@@ -1,8 +1,11 @@
+import io
 import logging
+import sys
 import unittest
+import unittest.mock
 from textwrap import dedent
 
-from snekbox.nsjail import MEM_MAX, NsJail
+from snekbox.nsjail import MEM_MAX, NsJail, READ_CHUNK_SIZE, OUTPUT_MAX
 
 
 class NsJailTests(unittest.TestCase):
@@ -183,3 +186,18 @@ class NsJailTests(unittest.TestCase):
 
         result = self.nsjail.python3(stdout_flood)
         self.assertEqual(result.returncode, 137)
+
+    def test_large_output_is_truncated(self):
+        chunk = "a" * READ_CHUNK_SIZE
+        expected_chunks = OUTPUT_MAX // sys.getsizeof(chunk) + 1
+
+        nsjail_subprocess = unittest.mock.Mock()
+
+        # 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)
+        self.assertEqual(output, chunk * expected_chunks)