Ensure assets are always bytes and encoded properly when serialized

coderunners/services.py

@@ -1,4 +1,3 @@
-import base64
 import gzip
 from pathlib import Path
 
@@ -125,12 +124,10 @@ def check(self) -> SubmissionResult:
                 test_results[-1].outputs = r.outputs[:max_len] if r.outputs else None
                 test_results[-1].errors = r.errors[:max_len] if r.errors else None
                 test_results[-1].output_files = {
-                    filename: content[:max_len]
-                    for filename, content in r.output_files.items()
+                    filename: content[:max_len] for filename, content in r.output_files.items()
                 } if r.output_files else None
                 test_results[-1].output_assets = {
-                    filename: base64.b64encode(content[:max_len]).decode('utf-8')
-                    for filename, content in r.output_assets.items()
+                    filename: content[:max_len] for filename, content in r.output_assets.items()
                 } if r.output_assets else None
 
             # Stop on failure

models.py

@@ -1,5 +1,5 @@
 import base64
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 
 from dataclasses_json import DataClassJsonMixin, LetterCase, Undefined, config
@@ -9,6 +9,18 @@ class DataClassJsonCamelMixIn(DataClassJsonMixin):
     dataclass_json_config = config(letter_case=LetterCase.CAMEL, undefined=Undefined.EXCLUDE)['dataclasses_json']
 
 
+def base64_to_bytes(data: dict[str, str | bytes] | None) -> dict[str, bytes] | None:
+    if data is not None and all(isinstance(content, str) for content in data.values()):
+        return {filename: base64.b64decode(content.encode('utf-8')) for filename, content in data.items()}
+    return data
+
+
+def bytes_to_base64(data: dict[str, bytes] | None) -> dict[str, str] | None:
+    if data is not None and all(isinstance(content, bytes) for content in data.values()):
+        return {filename: base64.b64encode(content).decode('utf-8') for filename, content in data.items()}
+    return data
+
+
 class Status(Enum):
     OK = 'Solved'
     WA = 'Wrong answer'
@@ -27,23 +39,14 @@ class TestCase(DataClassJsonCamelMixIn):
     target: str
     input_files: dict[str, str] | None = None               # mapping filename -> textual content
     target_files: dict[str, str] | None = None              # mapping filename -> textual content
-    input_assets: dict[str, str | bytes] | None = None      # mapping filename -> base64 encoded string or bytes
-    target_assets: dict[str, str | bytes] | None = None     # mapping filename -> base64 encoded string or bytes
-
-    def __post_init__(self):
-        """
-        Make sure that if input_assets or target_assets are provided (as base64 encoded strings), convert them to bytes
-        """
-        if self.input_assets is not None and all(isinstance(content, str) for content in self.input_assets.values()):
-            self.input_assets = {
-                filename: base64.b64decode(content.encode('utf-8'))
-                for filename, content in self.input_assets.items()
-            }
-        if self.target_assets is not None and all(isinstance(content, str) for content in self.target_assets.values()):
-            self.target_assets = {
-                filename: base64.b64decode(content.encode('utf-8'))
-                for filename, content in self.target_assets.items()
-            }
+    input_assets: dict[str, bytes] | None = field(          # mapping filename -> binary content
+        metadata=config(encoder=bytes_to_base64, decoder=base64_to_bytes),
+        default=None,
+    )
+    target_assets: dict[str, bytes] | None = field(         # mapping filename -> binary content
+        metadata=config(encoder=bytes_to_base64, decoder=base64_to_bytes),
+        default=None,
+    )
 
 
 @dataclass
@@ -108,7 +111,10 @@ class RunResult(DataClassJsonCamelMixIn):
     outputs: str | None = None
     errors: str | None = None
     output_files: dict[str, str] | None = None
-    output_assets: dict[str, str] | None = None
+    output_assets: dict[str, bytes] | None = field(
+        metadata=config(encoder=bytes_to_base64, decoder=base64_to_bytes),
+        default=None,
+    )
 
 
 @dataclass

sync/app.py

@@ -1,4 +1,3 @@
-import base64
 import json
 import os
 from pathlib import Path
@@ -38,17 +37,6 @@ def trigger_sync_s3_handler(event, context):
 
     tests = TestCase.schema().load(res['tests_truncated'], many=True)
     print('tests:', tests)
-
-    # Convert asset files to base64 strings
-    for test in tests:
-        test.input_assets = {
-            filename: base64.b64encode(content).decode('utf-8')
-            for filename, content in test.input_assets.items()
-        } if test.input_assets else None
-        test.target_assets = {
-            filename: base64.b64encode(content).decode('utf-8')
-            for filename, content in test.target_assets.items()
-        } if test.target_assets else None
     SummaryTable(dynamodb).write(problem, tests)
     print('Wrote to a summary table')
 
@@ -70,19 +58,6 @@ def sync_efs_handler(event, context):
 
     tests = zip2tests(zip_path)
     tests_truncated = truncate(tests, max_len=100)
-    print('tests_truncated:', tests_truncated)
-
-    # Convert asset files to base64 strings
-    for test in tests_truncated:
-        test.input_assets = {
-            filename: base64.b64encode(content).decode('utf-8')
-            for filename, content in test.input_assets.items()
-        } if test.input_assets else None
-        test.target_assets = {
-            filename: base64.b64encode(content).decode('utf-8')
-            for filename, content in test.target_assets.items()
-        } if test.target_assets else None
-    print('tests_truncated after conversion to str:', tests_truncated)
     tests = encrypt_tests(tests, encryption_key=encryption_key)
 
     with open(problem_file, 'wb') as f:

tests/integration/coderunners/test_files.py

@@ -1,4 +1,3 @@
-import base64
 from textwrap import dedent
 
 from bouncer.coderunners import CodeRunner
@@ -11,11 +10,9 @@ class TestFiles:
         TestCase(input='1', target='2', input_files={'hello.txt': 'hello'}, target_files={'hello.txt': 'hello'}),
         TestCase(input='2', target='4', input_files={'hello.txt': 'hello'}, target_files={'hello.txt': 'hello'}),
         TestCase(input='3', target='4', input_files={'hello.txt': 'hello'}, target_files={'res.txt': 'heyhey'}),
-        TestCase(input='4', target='5', input_files={'hello.txt': 'hello'}, target_files={'res.txt': 'heyhey'}),
+        TestCase(input='4', target='5', input_files={'hello.txt': 'hello'}, target_files={'res.txt': 'heyhey'},
+                 input_assets={'img.bmp': b'image!'}, target_assets={'res.bmp': b'Result!!!'}),
     ]
-    # Assets should be encoded into base64 strings
-    test_cases[-1].input_assets = {'img.bmp': base64.b64encode(b'image!').decode('utf-8')}
-    test_cases[-1].target_assets = {'res.bmp': base64.b64encode(b'Result!!!').decode('utf-8')}
 
     def test_no_file(self):
         request = SubmissionRequest(test_cases=self.test_cases, stop_on_first_fail=False, language='python', code={

tests/unit/__init__.py

@@ -0,0 +1,4 @@
+from models import TestCase, TestGroup
+
+TestCase.__test__ = False
+TestGroup.__test__ = False