Rewrite NsJail tests

* Fix SIGSEGV test
* Add embedded null byte test
* Return None for stderr when there's a ValueError

3 files changed, 80 insertions, 58 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 3bcc0a1..f82dcf0 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -5,6 +5,7 @@ import subprocess
 import sys
 import textwrap
 from pathlib import Path
+from subprocess import CompletedProcess
 from tempfile import NamedTemporaryFile
 from typing import List
 
@@ -92,7 +93,7 @@ class NsJail:
                 # Treat fatal as error.
                 log.error(msg)
 
-    def python3(self, code: str) -> subprocess.CompletedProcess:
+    def python3(self, code: str) -> CompletedProcess:
         """Execute Python 3 code in an isolated environment and return the completed process."""
         with NamedTemporaryFile() as nsj_log:
             args = (
@@ -130,7 +131,7 @@ class NsJail:
                     text=True
                 )
             except ValueError:
-                return subprocess.CompletedProcess(args, None, "ValueError: embedded null byte", "")
+                return CompletedProcess(args, None, "ValueError: embedded null byte", None)
 
             log_lines = nsj_log.read().decode("UTF-8").splitlines()
             if not log_lines and result.returncode == 255:
diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py
new file mode 100644
index 0000000..1184b87
--- /dev/null
+++ b/tests/test_nsjail.py
@@ -0,0 +1,77 @@
+import logging
+import unittest
+from textwrap import dedent
+
+from snekbox.nsjail import NsJail
+
+
+class NsJailTests(unittest.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        self.nsjail = NsJail()
+        self.logger = logging.getLogger("snekbox.nsjail")
+
+    def test_print_returns_0(self):
+        result = self.nsjail.python3("print('test')")
+        self.assertEqual(result.returncode, 0)
+        self.assertEqual(result.stdout, "test\n")
+        self.assertEqual(result.stderr, None)
+
+    def test_timeout_returns_137(self):
+        code = dedent("""
+            x = '*'
+            while True:
+                try:
+                    x = x * 99
+                except:
+                    continue
+        """).strip()
+
+        with self.assertLogs(self.logger) as log:
+            result = self.nsjail.python3(code)
+
+        self.assertEqual(result.returncode, 137)
+        self.assertEqual(result.stdout, "")
+        self.assertEqual(result.stderr, None)
+        self.assertIn("run time >= time limit", "\n".join(log.output))
+
+    def test_subprocess_resource_unavailable(self):
+        code = dedent("""
+            import subprocess
+            print(subprocess.check_output('kill -9 6', shell=True).decode())
+        """).strip()
+
+        result = self.nsjail.python3(code)
+        self.assertEqual(result.returncode, 1)
+        self.assertIn("Resource temporarily unavailable", result.stdout)
+        self.assertEqual(result.stderr, None)
+
+    def test_forkbomb_resource_unavailable(self):
+        code = dedent("""
+            import os
+            while 1:
+                os.fork()
+        """).strip()
+
+        result = self.nsjail.python3(code)
+        self.assertEqual(result.returncode, 1)
+        self.assertIn("Resource temporarily unavailable", result.stdout)
+        self.assertEqual(result.stderr, None)
+
+    def test_sigsegv_returns_139(self):  # In honour of Juan.
+        code = dedent("""
+            import ctypes
+            ctypes.string_at(0)
+        """).strip()
+
+        result = self.nsjail.python3(code)
+        self.assertEqual(result.returncode, 139)
+        self.assertEqual(result.stdout, "")
+        self.assertEqual(result.stderr, None)
+
+    def test_null_byte_value_error(self):
+        result = self.nsjail.python3("\0")
+        self.assertEqual(result.returncode, None)
+        self.assertEqual(result.stdout, "ValueError: embedded null byte")
+        self.assertEqual(result.stderr, None)
diff --git a/tests/test_snekbox.py b/tests/test_snekbox.py
deleted file mode 100644
index 46319d6..0000000
--- a/tests/test_snekbox.py
+++ /dev/null
@@ -1,56 +0,0 @@
-import unittest
-
-from snekbox.nsjail import NsJail
-
-nsjail = NsJail()
-
-
-class SnekTests(unittest.TestCase):
-    def test_nsjail(self):
-        result = nsjail.python3("print('test')")
-        self.assertEquals(result.strip(), "test")
-
-    # def test_memory_error(self):
-    #     code = ("x = "*"\n"
-    #             "while True:\n"
-    #             "    x = x * 99\n")
-    #     result = nsjail.python3(code)
-    #     self.assertEquals(result.strip(), "timed out or memory limit exceeded")
-
-    def test_timeout(self):
-        code = (
-            "x = '*'\n"
-            "while True:\n"
-            "    try:\n"
-            "        x = x * 99\n"
-            "    except:\n"
-            "        continue\n"
-        )
-
-        result = nsjail.python3(code)
-        self.assertEquals(result.strip(), "timed out or memory limit exceeded")
-
-    def test_kill(self):
-        code = ("import subprocess\n"
-                "print(subprocess.check_output('kill -9 6', shell=True).decode())")
-        result = nsjail.python3(code)
-        if "ModuleNotFoundError" in result.strip():
-            self.assertIn("ModuleNotFoundError", result.strip())
-        else:
-            self.assertIn("(PIDs left: 0)", result.strip())
-
-    def test_forkbomb(self):
-        code = ("import os\n"
-                "while 1:\n"
-                "    os.fork()")
-        result = nsjail.python3(code)
-        self.assertIn("Resource temporarily unavailable", result.strip())
-
-    def test_juan_golf(self):  # in honour of Juan
-        code = ("func = lambda: None\n"
-                "CodeType = type(func.__code__)\n"
-                "bytecode = CodeType(0,1,0,0,0,b'',(),(),(),'','',1,b'')\n"
-                "exec(bytecode)")
-
-        result = nsjail.python3(code)
-        self.assertEquals("unknown error, code: 111", result.strip())