Refactor writing of input files into a separate function

3 files changed, 31 insertions, 18 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 4e86f33..b1d6547 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -12,7 +12,7 @@ from google.protobuf import text_format
 from snekbox import DEBUG, limits
 from snekbox.config_pb2 import NsJailConfig
 from snekbox.limits.timed import time_limit
-from snekbox.process import EvalResult
+from snekbox.result import EvalError, EvalResult
 from snekbox.snekio import FileAttachment, MemFS
 from snekbox.snekio.filesystem import Size
 from snekbox.utils.iter import iter_lstrip
@@ -199,6 +199,24 @@ class NsJail:
             *iter_lstrip(py_args),
         ]
 
+    def _write_files(self, home: Path, files: Iterable[FileAttachment]) -> dict[Path, float]:
+        files_written = {}
+        for file in files:
+            try:
+                f_path = file.save_to(home)
+                # Allow file to be writable
+                f_path.chmod(0o777)
+                # Save the written at time to later check if it was modified
+                files_written[f_path] = f_path.stat().st_mtime
+                log.info(f"Created file at {(home / file.path)!r}.")
+            except OSError as e:
+                log.info(f"Failed to create file at {(home / file.path)!r}.", exc_info=e)
+                raise EvalError(
+                    f"{e.__class__.__name__}: Failed to create file '{file.path}'."
+                ) from e
+
+        return files_written
+
     def python3(
         self,
         py_args: Iterable[str],
@@ -220,21 +238,10 @@ class NsJail:
         ) as fs:
             args = self._build_args(py_args, nsjail_args, nsj_log.name, str(fs.home))
 
-            # Write provided files if any
-            files_written: dict[Path, float] = {}
-            for file in files:
-                try:
-                    f_path = file.save_to(fs.home)
-                    # Allow file to be writable
-                    f_path.chmod(0o777)
-                    # Save the written at time to later check if it was modified
-                    files_written[f_path] = f_path.stat().st_mtime
-                    log.info(f"Created file at {(fs.home / file.path)!r}.")
-                except OSError as e:
-                    log.info(f"Failed to create file at {(fs.home / file.path)!r}.", exc_info=e)
-                    return EvalResult(
-                        args, None, f"{e.__class__.__name__}: Failed to create file '{file.path}'."
-                    )
+            try:
+                files_written = self._write_files(fs.home, files)
+            except EvalError as e:
+                return EvalResult(args, None, str(e))
 
             msg = "Executing code..."
             if DEBUG:
diff --git a/snekbox/process.py b/snekbox/result.py
index 552b91a..e11d0e4 100644
--- a/snekbox/process.py
+++ b/snekbox/result.py
@@ -1,4 +1,4 @@
-"""Utilities for process management."""
+"""Types for representing the result of an evaluation job."""
 from collections.abc import Sequence
 from os import PathLike
 from subprocess import CompletedProcess
@@ -6,6 +6,8 @@ from typing import TypeVar
 
 from snekbox.snekio import FileAttachment
 
+__all__ = ("EvalError", "EvalResult")
+
 _T = TypeVar("_T")
 ArgType = (
     str
@@ -16,6 +18,10 @@ ArgType = (
 )
 
 
+class EvalError(RuntimeError):
+    """An error that occurred during evaluation."""
+
+
 class EvalResult(CompletedProcess[_T]):
     """An evaluation job that has finished running."""
 
diff --git a/tests/api/__init__.py b/tests/api/__init__.py
index 5f20faf..c84080e 100644
--- a/tests/api/__init__.py
+++ b/tests/api/__init__.py
@@ -4,7 +4,7 @@ from unittest import mock
 from falcon import testing
 
 from snekbox.api import SnekAPI
-from snekbox.process import EvalResult
+from snekbox.result import EvalResult
 
 
 class SnekAPITestCase(testing.TestCase):