1 files changed, 90 insertions, 0 deletions
diff --git a/resources/unittest_template.py b/resources/unittest_template.py
new file mode 100644
index 0000000..2410278
--- /dev/null
+++ b/resources/unittest_template.py
@@ -0,0 +1,90 @@
+# flake8: noqa
+"""This template is used inside snekbox to evaluate and test user code."""
+import ast
+import base64
+import io
+import os
+import sys
+import traceback
+import unittest
+from itertools import chain
+from types import ModuleType, SimpleNamespace
+from typing import NoReturn
+from unittest import mock
+
+### USER CODE
+
+
+class RunnerTestCase(unittest.TestCase):
+### UNIT CODE
+
+
+def _exit_sandbox(code: int) -> NoReturn:
+    """
+    Exit the sandbox by printing the result to the actual stdout and exit with the provided code.
+
+    Codes:
+    - 0: Executed with success
+    - 5: Syntax error while parsing user code
+    - 6: Uncaught exception while loading user code
+    - 99: Internal error
+
+    137 can also be generated by NsJail when killing the process.
+    """
+    print(RESULT.getvalue(), file=ORIGINAL_STDOUT, end="")
+    sys.exit(code)
+
+
+def _load_user_module() -> ModuleType:
+    """Load the user code into a new module and return it."""
+    code = base64.b64decode(USER_CODE).decode("utf8")
+    try:
+        ast.parse(code, "<input>")
+    except SyntaxError:
+        RESULT.write("".join(traceback.format_exception(*sys.exc_info(), limit=0)))
+        _exit_sandbox(5)
+
+    _module = ModuleType("module")
+    exec(code, _module.__dict__)
+
+    return _module
+
+
+def _main() -> None:
+    suite = unittest.defaultTestLoader.loadTestsFromTestCase(RunnerTestCase)
+    result = suite.run(unittest.TestResult())
+
+    RESULT.write(str(int(result.wasSuccessful())))
+
+    if not result.wasSuccessful():
+        RESULT.write(
+            ";".join(chain(
+                (error[0]._testMethodName.lstrip("test_") for error in result.errors),
+                (failure[0]._testMethodName.lstrip("test_") for failure in result.failures)
+            ))
+        )
+
+    _exit_sandbox(0)
+
+
+try:
+    # Fake file object not writing anything
+    DEVNULL = SimpleNamespace(write=lambda *_: None, flush=lambda *_: None)
+
+    RESULT = io.StringIO()
+    ORIGINAL_STDOUT = sys.stdout
+
+    # stdout/err is patched in order to control what is outputted by the runner
+    sys.stdout = DEVNULL
+    sys.stderr = DEVNULL
+    
+    # Load the user code as a global module variable
+    try:
+        module = _load_user_module()
+    except Exception:
+        RESULT.write("Uncaught exception while loading user code.")
+        _exit_sandbox(6)
+    _main()
+except Exception:
+    RESULT.write("Uncaught exception inside runner.")
+    _exit_sandbox(99)