Add memory limit, defaulted to 1GB

client/cli/ok.py

@@ -43,6 +43,7 @@
 
 Visit https://okpy.org to view your backups and submissions.
 """
+import resource
 
 from client import exceptions as ex
 from client.api import assignment
@@ -105,6 +106,8 @@ def parse_input(command_input=None):
                         help="submit composition revision")
     testing.add_argument('--timeout', type=int, default=10,
                         help="set the timeout duration (in seconds) for running tests")
+    testing.add_argument('--max-memory', type=int, default=1024,
+                        help="set the maximal memory usage (in mebibytes) for running tests")
     testing.add_argument('-cov', '--coverage', action='store_true',
                         help="get suggestions on what lines to add tests for")
 
@@ -185,6 +188,9 @@ def main():
     log.setLevel(logging.DEBUG if args.debug else logging.ERROR)
     log.debug(args)
 
+    _, hard = resource.getrlimit(resource.RLIMIT_AS)
+    resource.setrlimit(resource.RLIMIT_AS, (2 ** 20 * args.max_memory, hard))
+
     # Checking user's Python bit version
     bit_v = (8 * struct.calcsize("P"))
     log.debug("Python {} ({}bit)".format(sys.version, bit_v))

tests/end_to_end/smoke_test.py

@@ -2,20 +2,18 @@
 
 from tests.end_to_end.end_to_end_test import EndToEndTest
 
+TOO_MUCH_MEMORY_TEST = """
+def f():
+    '''
+    >>> f()
+    2
+    '''
+    return 0 * len(str(list(range(10 ** 7)))) + 2
+"""
 
 class SmokeTest(EndToEndTest):
 
-    def testVersion(self):
-        stdout, stderr = self.run_ok("--version")
-        self.assertEqual(stderr, "")
-        self.assertRegex(stdout, "^okpy==.*")
-
-    def testUpdate(self):
-        stdout, stderr = self.run_ok("--update")
-        self.assertEqual(stderr, "")
-        self.assertRegex(stdout, "Current version: v[0-9.]+\nOK is up to date")
-
-    def testRunNoArgument(self):
+    def add_test_ok(self):
         self.add_file("test.ok", json.dumps(
             {
                 "name": "Test Assignment",
@@ -37,6 +35,30 @@ def testRunNoArgument(self):
                 ]
             }
         ))
+
+    def testVersion(self):
+        stdout, stderr = self.run_ok("--version")
+        self.assertEqual(stderr, "")
+        self.assertRegex(stdout, "^okpy==.*")
+
+    def testUpdate(self):
+        stdout, stderr = self.run_ok("--update")
+        self.assertEqual(stderr, "")
+        self.assertRegex(stdout, "Current version: v[0-9.]+\nOK is up to date")
+
+    def testRunNoArgument(self):
+        self.add_test_ok()
         stdout, stderr = self.run_ok("--local")
         self.assertEqual(stderr, "")
         self.assertRegex(stdout, "0 test cases passed! No cases failed")
+
+    def testRunTooMuchMemory(self):
+        self.add_test_ok()
+        self.add_file("test.py", TOO_MUCH_MEMORY_TEST)
+        stdout, stderr = self.run_ok("--local", "-q", "f", "--max-memory", "100")
+        self.assertEqual(stderr, "")
+        self.assertRegex(stdout, "MemoryError")
+        self.assertRegex(stdout, "0 test cases passed before encountering first failed test case")
+        stdout, stderr = self.run_ok("--local", "-q", "f", "--max-memory", "1000")
+        self.assertEqual(stderr, "")
+        self.assertRegex(stdout, "1 test cases passed! No cases failed")