Change Size to IntEnum, change usage in tests

3 files changed, 11 insertions, 9 deletions

diff --git a/snekbox/filesystem.py b/snekbox/filesystem.py
index 0cdfbf5..9f0634a 100644
--- a/snekbox/filesystem.py
+++ b/snekbox/filesystem.py
@@ -20,13 +20,13 @@ libc.mount.argtypes = (
 libc.umount2.argtypes = (ctypes.c_char_p, ctypes.c_int)
 
 
-class Size(int):
-    """Size in bytes."""
+class Size(IntEnum):
+    """Size multipliers for bytes."""
 
-    @classmethod
-    def from_mb(cls, mb: int) -> Size:
-        """Create a Size from mebibytes."""
-        return cls(mb * 1024 * 1024)
+    KB = 1024
+    MB = 1024**2
+    GB = 1024**3
+    TB = 1024**4
 
 
 class UnmountFlags(IntEnum):
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 0dd49e5..887bc9a 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -10,6 +10,7 @@ from google.protobuf import text_format
 
 from snekbox import DEBUG, utils
 from snekbox.config_pb2 import NsJailConfig
+from snekbox.filesystem import Size
 from snekbox.memfs import MemFS
 from snekbox.process import EvalResult
 from snekbox.snekio import FileAttachment
@@ -50,7 +51,7 @@ class NsJail:
         config_path: str = "./config/snekbox.cfg",
         max_output_size: int = 1_000_000,
         read_chunk_size: int = 10_000,
-        memfs_instance_size: int = 48 * 1024 * 1024,
+        memfs_instance_size: int = 48 * Size.MB,
         files_limit: int | None = 100,
         files_timeout: float | None = 8,
         files_pattern: str = "**/*",
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(