1 files changed, 26 insertions, 0 deletions
diff --git a/backend/models/question.py b/backend/models/question.py
index 7daeb5a..5a1334a 100644
--- a/backend/models/question.py
+++ b/backend/models/question.py
@@ -4,6 +4,28 @@ from pydantic import BaseModel, Field, root_validator, validator
 
 from backend.constants import QUESTION_TYPES, REQUIRED_QUESTION_TYPE_DATA
 
+_TESTS_TYPE = t.Union[t.Dict[str, str], int]
+
+
+class Unittests(BaseModel):
+    """Schema model for unittest suites in code questions."""
+    allow_failure: bool = False
+    tests: _TESTS_TYPE
+
+    @validator("tests")
+    def validate_tests(cls, value: _TESTS_TYPE) -> _TESTS_TYPE:
+        """Confirm that at least one test exists in a test suite."""
+        if isinstance(value, dict) and len(value.keys()) == 0:
+            raise ValueError("Must have at least one test in a test suite.")
+
+        return value
+
+
+class CodeQuestion(BaseModel):
+    """Schema model for questions of type `code`."""
+    language: str
+    unittests: t.Optional[Unittests]
+
 
 class Question(BaseModel):
     """Schema model for form question."""
@@ -49,4 +71,8 @@ class Question(BaseModel):
                     f"got {type(value['data'][key]).__name__} instead."
                 )
 
+            # Validate unittest options
+            if value.get("type").lower() == "code":
+                value["data"] = CodeQuestion(**value.get("data")).dict()
+
         return value