Merge #132 - fix entrypoint and support Python args

3 files changed, 119 insertions, 9 deletions

diff --git a/DEVELOPING.md b/DEVELOPING.md
index 0dbf4cc..ef758bc 100644
--- a/DEVELOPING.md
+++ b/DEVELOPING.md
@@ -85,7 +85,7 @@ pipenv run devsh -c 'echo hello'
 NsJail can be invoked in a more direct manner that does not require using a web server or its API. See `python -m snekbox --help`. Example usage:
 
 ```bash
-python -m snekbox 'print("hello world!")' --time_limit 0
+python -m snekbox 'print("hello world!")' --time_limit 0 --- -m timeit
 ```
 
 With this command, NsJail uses the same configuration normally used through the web API. It also has an alias, `pipenv run eval`.
diff --git a/snekbox/__main__.py b/snekbox/__main__.py
index a63fd3c..704ec9d 100644
--- a/snekbox/__main__.py
+++ b/snekbox/__main__.py
@@ -1,28 +1,48 @@
 import argparse
+import sys
 
 from snekbox.nsjail import NsJail
 
 
 def parse_args() -> argparse.Namespace:
     """Parse the command-line arguments and return the populated namespace."""
-    parser = argparse.ArgumentParser(prog="snekbox", usage="%(prog)s code [nsjail_args ...]")
+    parser = argparse.ArgumentParser(
+        prog="snekbox",
+        usage="%(prog)s [-h] code [nsjail_args ...] [--- py_args ...]",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
     parser.add_argument("code", help="the Python code to evaluate")
-    parser.add_argument("nsjail_args", nargs="?", help="override configured NsJail options")
-
-    # nsjail_args is just a dummy for documentation purposes.
-    # Its actual value comes from all the unknown arguments.
+    parser.add_argument(
+        "nsjail_args", nargs="?", default=[], help="override configured NsJail options"
+    )
+    parser.add_argument(
+        "py_args", nargs="?", default=["-c"], help="arguments to pass to the Python process"
+    )
+
+    # nsjail_args and py_args are just dummies for documentation purposes.
+    # Their actual values come from all the unknown arguments.
     # There doesn't seem to be a better solution with argparse.
     args, unknown = parser.parse_known_args()
-    args.nsjail_args = unknown
+    try:
+        # Can't use double dash because that has special semantics for argparse already.
+        split = unknown.index("---")
+        args.nsjail_args = unknown[:split]
+        args.py_args = unknown[split + 1:]
+    except ValueError:
+        args.nsjail_args = unknown
+
     return args
 
 
 def main() -> None:
     """Evaluate Python code through NsJail."""
     args = parse_args()
-    result = NsJail().python3(args.code, *args.nsjail_args)
+    result = NsJail().python3(args.code, nsjail_args=args.nsjail_args, py_args=args.py_args)
     print(result.stdout)
 
+    if result.returncode != 0:
+        sys.exit(result.returncode)
+
 
-if __name__ == "__main__":
+if __name__ == "__main__":  # pragma: no cover
     main()
diff --git a/tests/test_main.py b/tests/test_main.py
new file mode 100644
index 0000000..5b5bfcb
--- /dev/null
+++ b/tests/test_main.py
@@ -0,0 +1,90 @@
+import ast
+import contextlib
+import io
+import logging
+import unittest
+from argparse import Namespace
+from unittest.mock import patch
+
+import snekbox.__main__ as snekbox_main
+
+
+class ArgParseTests(unittest.TestCase):
+    def test_parse_args(self):
+        subtests = (
+            (
+                ["", "code"],
+                Namespace(code="code", nsjail_args=[], py_args=["-c"])
+            ),
+            (
+                ["", "code", "--time_limit", "0"],
+                Namespace(code="code", nsjail_args=["--time_limit", "0"], py_args=["-c"])
+            ),
+            (
+                ["", "code", "---", "-m", "timeit"],
+                Namespace(code="code", nsjail_args=[], py_args=["-m", "timeit"])
+            ),
+            (
+                ["", "code", "--time_limit", "0", "---", "-m", "timeit"],
+                Namespace(code="code", nsjail_args=["--time_limit", "0"], py_args=["-m", "timeit"])
+            ),
+            (
+                ["", "code", "--time_limit", "0", "---"],
+                Namespace(code="code", nsjail_args=["--time_limit", "0"], py_args=[])
+            ),
+            (
+                ["", "code", "---"],
+                Namespace(code="code", nsjail_args=[], py_args=[])
+            )
+        )
+
+        for argv, expected in subtests:
+            with self.subTest(argv=argv, expected=expected), patch("sys.argv", argv):
+                args = snekbox_main.parse_args()
+                self.assertEqual(args, expected)
+
+    @patch("sys.argv", [""])
+    def test_parse_args_code_missing_exits(self):
+        with self.assertRaises(SystemExit) as cm:
+            with contextlib.redirect_stderr(io.StringIO()) as stderr:
+                snekbox_main.parse_args()
+
+        self.assertEqual(cm.exception.code, 2)
+        self.assertIn("the following arguments are required: code", stderr.getvalue())
+
+
+class EntrypointTests(unittest.TestCase):
+    """Integration tests of the CLI entrypoint."""
+
+    def setUp(self):
+        logging.getLogger("snekbox.nsjail").setLevel(logging.WARNING)
+
+    @patch("sys.argv", ["", "print('hello'); import sys; print('error', file=sys.stderr)"])
+    def test_main_prints_stdout(self):
+        """Should print the stdout of the subprocess followed by its stderr."""
+        with contextlib.redirect_stdout(io.StringIO()) as stdout:
+            snekbox_main.main()
+
+        self.assertEqual(stdout.getvalue(), "hello\nerror\n\n")
+
+    @patch("sys.argv", ["", "import sys; sys.exit(22)"])
+    def test_main_exits_with_returncode(self):
+        """Should exit with the subprocess's returncode if it's non-zero."""
+        with self.assertRaises(SystemExit) as cm:
+            snekbox_main.main()
+
+        self.assertEqual(cm.exception.code, 22)
+
+    def test_main_forwards_args(self):
+        """Should forward NsJail args to NsJail and Python args to the Python subprocess."""
+        code = "import sys, time; print(sys.orig_argv); time.sleep(2)"
+        py_args = ["-R", "-dc"]
+        args = ["", code, "--time_limit", "1", "---", *py_args]
+
+        with patch("sys.argv", args), self.assertRaises(SystemExit) as cm:
+            with contextlib.redirect_stdout(io.StringIO()) as stdout:
+                snekbox_main.main()
+
+        orig_argv = ast.literal_eval(stdout.getvalue().strip())
+        self.assertListEqual([*py_args, code], orig_argv[-3:])
+        self.assertEqual(cm.exception.code, 137, "The time_limit NsJail arg was not respected.")