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.

1 files changed, 19 insertions, 1 deletions

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)