Refactor output files in `output` dir

5 files changed, 54 insertions, 20 deletions

diff --git a/snekbox/memfs.py b/snekbox/memfs.py
index 16f5cfb..a39b690 100644
--- a/snekbox/memfs.py
+++ b/snekbox/memfs.py
@@ -62,7 +62,12 @@ class MemFS:
     @property
     def home(self) -> Path:
         """Path to home directory."""
-        return Path(self.path, "home")
+        return self.path / "home"
+
+    @property
+    def output(self) -> Path:
+        """Path to output directory."""
+        return self.home / "output"
 
     def __enter__(self) -> MemFS:
         """Mounts a new tempfs, returns self."""
@@ -79,7 +84,8 @@ class MemFS:
         else:
             raise RuntimeError("Failed to generate a unique tempdir name in 10 attempts")
 
-        self.mkdir("home")
+        self.mkdir(self.home)
+        self.mkdir(self.output)
         return self
 
     def __exit__(
@@ -98,17 +104,17 @@ class MemFS:
         return folder
 
     def attachments(
-        self, max_count: int, pattern: str = "output*"
+        self, max_count: int, pattern: str = "**/*"
     ) -> Generator[FileAttachment, None, None]:
         """Return a list of attachments in the tempdir."""
         count = 0
-        for file in self.home.glob(pattern):
+        for file in self.output.rglob(pattern):
             if count > max_count:
                 log.info(f"Max attachments {max_count} reached, skipping remaining files")
                 break
             if file.is_file():
                 count += 1
-                yield FileAttachment.from_path(file)
+                yield FileAttachment.from_path(file, relative_to=self.output)
 
     def cleanup(self) -> None:
         """Unmounts tmpfs."""
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 88af3cf..4538287 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -53,8 +53,21 @@ class NsJail:
         memfs_instance_size: int = 48 * 1024 * 1024,
         files_limit: int = 100,
         files_timeout: float = 15,
-        files_pattern: str = "output*",
+        files_pattern: str = "**/*",
     ):
+        """
+        Initialize NsJail.
+
+        Args:
+            nsjail_path: Path to the NsJail binary.
+            config_path: Path to the NsJail configuration file.
+            max_output_size: Maximum size of the output in bytes.
+            read_chunk_size: Size of the read buffer in bytes.
+            memfs_instance_size: Size of the tmpfs instance in bytes.
+            files_limit: Maximum number of files to parse for attach.
+            files_timeout: Maximum time in seconds to wait for files to be written / read.
+            files_pattern: Pattern to match files to attach.
+        """
         self.nsjail_path = nsjail_path
         self.config_path = config_path
         self.max_output_size = max_output_size
@@ -178,7 +191,6 @@ class NsJail:
             )
 
         with NamedTemporaryFile() as nsj_log, MemFS(self.memfs_instance_size) as fs:
-            # Add the temp dir to be mounted as cwd
             nsjail_args = (
                 # Set fslimit to unlimited, cannot be set in cfg
                 # due to upstream protobuf parsing issue
@@ -204,7 +216,7 @@ class NsJail:
                 *iter_lstrip(py_args),
             ]
 
-            # Write files if any
+            # Write provided files if any
             for file in files:
                 file.save_to(fs.home)
                 log.info(f"Created file at {(fs.home / file.path)!r}.")
diff --git a/snekbox/snekio.py b/snekbox/snekio.py
index bc71501..7852301 100644
--- a/snekbox/snekio.py
+++ b/snekbox/snekio.py
@@ -57,9 +57,16 @@ class FileAttachment(Generic[T]):
         return cls(path, content)
 
     @classmethod
-    def from_path(cls, file: Path) -> FileAttachment[bytes]:
-        """Create an attachment from a file path."""
-        return cls(file.name, file.read_bytes())
+    def from_path(cls, file: Path, relative_to: Path | None = None) -> FileAttachment[bytes]:
+        """
+        Create an attachment from a file path.
+
+        Args:
+            file: The file to attach.
+            relative_to: The root for the path name.
+        """
+        path = file.relative_to(relative_to) if relative_to else file
+        return cls(str(path), file.read_bytes())
 
     @property
     def size(self) -> int:
diff --git a/tests/test_integration.py b/tests/test_integration.py
index 086abab..ba0d9b5 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -62,11 +62,15 @@ class IntegrationTests(unittest.TestCase):
                         "path": "main.py",
                         "content": b64encode_code(
                             """
+                            from pathlib import Path
                             from mod import lib
                             print(lib.var)
 
-                            with open('output.txt', 'w') as f:
-                                f.write('file write test')
+                            with open('output/test.txt', 'w') as f:
+                                f.write('test 1')
+
+                            Path('output/dir').mkdir()
+                            Path('output/dir/test2.txt').write_text('test 2')
                             """
                         ),
                     },
@@ -80,10 +84,15 @@ class IntegrationTests(unittest.TestCase):
                 "returncode": 0,
                 "files": [
                     {
-                        "path": "output.txt",
-                        "size": len("file write test"),
-                        "content": b64encode_code("file write test"),
-                    }
+                        "path": "dir/test2.txt",
+                        "size": len("test 2"),
+                        "content": b64encode_code("test 2"),
+                    },
+                    {
+                        "path": "test.txt",
+                        "size": len("test 1"),
+                        "content": b64encode_code("test 1"),
+                    },
                 ],
             }
 
diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py
index da2afea..d63180d 100644
--- a/tests/test_nsjail.py
+++ b/tests/test_nsjail.py
@@ -190,16 +190,16 @@ class NsJailTests(unittest.TestCase):
             data = "a" * 1024
             size = 32 * 1024 * 1024
 
-            with open("src", "w") as f:
+            with open("output/file", "w") as f:
                 for _ in range((size // 1024) - 5):
                     f.write(data)
 
             for i in range(100):
-                os.symlink("src", f"output{i}")
+                os.symlink("file", f"output/file{i}")
             """
         ).strip()
 
-        nsjail = NsJail(memfs_instance_size=48 * 1024 * 1024, files_timeout=1)
+        nsjail = NsJail(memfs_instance_size=32 * 1024 * 1024, files_timeout=1)
         result = nsjail.python3(["-c", code])
         self.assertEqual(result.returncode, None)
         self.assertEqual(