Add a test for changing limits.

codejail/jail_code.py

@@ -82,11 +82,14 @@ def set_limit(limit_name, value):
     value to use for that limit.  The type, meaning, default, and range of
     accepted values depend on `limit_name`.
 
-    There is currently only one limit defined:
+    These limits are available:
 
         * `"CPU"`: the maximum number of CPU seconds the jailed code can use.
             The value is an integer, defaulting to 1.
 
+        * `"VMEM"`: the total virtual memory available to the jailed code, in
+            bytes.  The default is 30 Mb.
+
     Limits are process-wide, and will affect all future calls to jail_code.
 
     """

codejail/tests/test_jailpy.py

@@ -5,7 +5,7 @@ import textwrap
 import unittest
 from nose.plugins.skip import SkipTest
 
-from codejail.jail_code import jail_code, is_configured
+from codejail.jail_code import jail_code, is_configured, set_limit, LIMITS
 
 dedent = textwrap.dedent
 
@@ -84,6 +84,7 @@ class TestFeatures(JailCodeHelpers, unittest.TestCase):
 
 class TestLimits(JailCodeHelpers, unittest.TestCase):
     def test_cant_use_too_much_memory(self):
+        # Default is 30Mb, try to get a 50Mb object.
         res = jailpy(code="print len(bytearray(50000000))")
         self.assertEqual(res.stdout, "")
         self.assertNotEqual(res.status, 0)
@@ -138,6 +139,28 @@ class TestLimits(JailCodeHelpers, unittest.TestCase):
         self.assertIn("OSError", res.stderr)
 
 
+class TestChangingLimits(JailCodeHelpers, unittest.TestCase):
+    def setUp(self):
+        super(TestChangingLimits, self).setUp()
+        self.old_vm = LIMITS['VMEM']
+
+    def tearDown(self):
+        set_limit('VMEM', self.old_vm)
+        super(TestChangingLimits, self).tearDown()
+
+    def test_changing_vmem_limit(self):
+        # This will fail (default is 30Mb)
+        res = jailpy(code="print len(bytearray(50000000))")
+        self.assertEqual(res.stdout, "")
+        self.assertNotEqual(res.status, 0)
+
+        # Up the limit, it will succeed
+        set_limit('VMEM', 600000000)
+        res = jailpy(code="print len(bytearray(50000000))")
+        self.assertEqual(res.stdout, "50000000\n")
+        self.assertEqual(res.status, 0)
+
+
 class TestMalware(JailCodeHelpers, unittest.TestCase):
     def test_crash_cpython(self):
         # http://nedbatchelder.com/blog/201206/eval_really_is_dangerous.html