Change Size to IntEnum, change usage in tests

1 files changed, 3 insertions, 2 deletions

diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py
index 839d3ec..927b8c0 100644
--- a/tests/test_nsjail.py
+++ b/tests/test_nsjail.py
@@ -9,6 +9,7 @@ from itertools import product
 from pathlib import Path
 from textwrap import dedent
 
+from snekbox.filesystem import Size
 from snekbox.nsjail import NsJail
 from snekbox.snekio import FileAttachment
 
@@ -18,7 +19,7 @@ class NsJailTests(unittest.TestCase):
         super().setUp()
 
         # Specify lower limits for unit tests to complete within time limits
-        self.nsjail = NsJail(memfs_instance_size=2 * 1024 * 1024)
+        self.nsjail = NsJail(memfs_instance_size=2 * Size.MB)
         self.logger = logging.getLogger("snekbox.nsjail")
         self.logger.setLevel(logging.WARNING)
 
@@ -199,7 +200,7 @@ class NsJailTests(unittest.TestCase):
             """
         ).strip()
 
-        nsjail = NsJail(memfs_instance_size=32 * 1024 * 1024, files_timeout=1)
+        nsjail = NsJail(memfs_instance_size=32 * Size.MB, files_timeout=1)
         result = nsjail.python3(["-c", code])
         self.assertEqual(result.returncode, None)
         self.assertEqual(