Add file send integration tests

1 files changed, 54 insertions, 3 deletions

diff --git a/tests/test_integration.py b/tests/test_integration.py
index eba5e60..086abab 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -1,14 +1,25 @@
 import json
 import unittest
 import urllib.request
+from base64 import b64encode
 from multiprocessing.dummy import Pool
+from textwrap import dedent
 
 from tests.gunicorn_utils import run_gunicorn
 
 
-def run_code_in_snekbox(code: str) -> tuple[str, int]:
+def b64encode_code(data: str):
+    data = dedent(data).strip()
+    return b64encode(data.encode()).decode("ascii")
+
+
+def snekbox_run_code(code: str) -> tuple[str, int]:
     body = {"args": ["-c", code]}
-    json_data = json.dumps(body).encode("utf-8")
+    return snekbox_request(body)
+
+
+def snekbox_request(content: dict) -> tuple[str, int]:
+    json_data = json.dumps(content).encode("utf-8")
 
     req = urllib.request.Request("http://localhost:8060/eval")
     req.add_header("Content-Type", "application/json; charset=utf-8")
@@ -34,9 +45,49 @@ class IntegrationTests(unittest.TestCase):
 
             args = [code] * processes
             with Pool(processes) as p:
-                results = p.map(run_code_in_snekbox, args)
+                results = p.map(snekbox_run_code, args)
 
             responses, statuses = zip(*results)
 
             self.assertTrue(all(status == 200 for status in statuses))
             self.assertTrue(all(json.loads(response)["returncode"] == 0 for response in responses))
+
+    def test_files_send_receive(self):
+        """Test sending and receiving files to snekbox."""
+        with run_gunicorn():
+            request = {
+                "args": ["main.py"],
+                "files": [
+                    {
+                        "path": "main.py",
+                        "content": b64encode_code(
+                            """
+                            from mod import lib
+                            print(lib.var)
+
+                            with open('output.txt', 'w') as f:
+                                f.write('file write test')
+                            """
+                        ),
+                    },
+                    {"path": "mod/__init__.py"},
+                    {"path": "mod/lib.py", "content": b64encode_code("var = 'hello'")},
+                ],
+            }
+
+            expected = {
+                "stdout": "hello\n",
+                "returncode": 0,
+                "files": [
+                    {
+                        "path": "output.txt",
+                        "size": len("file write test"),
+                        "content": b64encode_code("file write test"),
+                    }
+                ],
+            }
+
+            response, status = snekbox_request(request)
+
+            self.assertEqual(200, status)
+            self.assertEqual(expected, json.loads(response))