feat: add stdout for sample tests (#46)

Author: weebao

app/services/execution/docker.py

@@ -75,13 +75,15 @@ def get_run_commands(self, lang: str, file_path: str) -> list:
         raise ValueError(f"Unsupported language: {lang}")
 
     def run_container(
-        self, lang: str, file_path: str, difficulty: str
+        self, lang: str, file_path: str, difficulty: str, line_offset: int
     ) -> ExecutionResult:
         """
         Run the code in a Docker container.
 
+        :param lang: The language of the code.
         :param file_path: The path to the file to run.
         :param difficulty: The difficulty of the problem.
+        :param line_offset: The line offset for error logs.
         :return: The result of the execution.
         """
         # Get the memory and time limits for the difficulty level.
@@ -131,12 +133,14 @@ def run_container(
                     return ExecutionResult(
                         success=False,
                         message="Runtime Error Detected\n" + self.last_logs.strip(),
+                        line_offset=line_offset,
                     )
 
                 if self.last_stderr.strip():
                     return ExecutionResult(
                         success=False,
                         message="Runtime Error Detected\n" + self.last_stderr.strip(),
+                        line_offset=line_offset,
                     )
 
                 base_name = os.path.basename(file_path).split(".")[0]
@@ -150,6 +154,7 @@ def run_container(
                         success=True,
                         test_results=execution_data["hidden_results"]["test_results"],
                         sample_results=execution_data["sample_results"]["test_results"],
+                        line_offset=line_offset,
                     )
                 else:
                     return ExecutionResult(

app/services/execution/service.py

@@ -64,7 +64,7 @@ async def execute_code(
 
             # Create a temporary file to store the test runner file.
             with tempfile.NamedTemporaryFile(
-                mode="w", suffix=gen.get_file_extension(lang), delete=False
+                mode="w", suffix=gen.get_file_extension(), delete=False
             ) as f:
                 # Test data are pairs of test cases and its expected results.
                 test_data = [
@@ -83,7 +83,9 @@ async def execute_code(
                 f.write(file_content)
                 file_path = f.name
             try:
-                result = self.docker.run_container(lang, file_path, difficulty)
+                result = self.docker.run_container(
+                    lang, file_path, difficulty, gen.get_line_offset()
+                )
 
                 # If all tests passed, get runtime analysis
                 if result.all_cleared() and not settings.TESTING:

app/services/execution/templates.py

@@ -1,5 +1,7 @@
 PYTHON_TEMPLATE = r"""import json
 import traceback
+import io
+import sys
 from typing import *
 
 {code}
@@ -10,23 +12,26 @@ def __init__(
         expected: str,
         passed: bool,
         output: Any = None,
+        logs: str = None,
         error: str = None,
         input: str = None,
     ):
         self.expected = expected
         self.output = str(output) if output is not None else None
         self.passed = passed
+        self.logs = logs
         self.error = error
         self.input = input
         
-    def to_dict(self, include_input: bool = True):
+    def to_dict(self, is_sample: bool = True):
         result = {{
             "expected": self.expected,
             "output": self.output,
             "passed": self.passed,
             "error": self.error,
         }}
-        if include_input:
+        if is_sample:
+            result["logs"] = self.logs
             result["input"] = self.input
         return result
     
@@ -47,22 +52,30 @@ def run_tests(solution, method_name, test_data, is_sample: bool = False):
     results = []
     
     for test in test_data:
+        old_stdout = sys.stdout
+        new_stdout = io.StringIO()
+        sys.stdout = new_stdout
+
         try:
             result = eval(f"solution.{{format_test_data(method_name, test['input'])}}")
             passed = compare_results(result, test['expected'])
             results.append(TestResult(
                 expected=test['expected'],
                 output=result,
                 passed=passed,
+                logs=new_stdout.getvalue(),
                 input=test['input'],
-            ).to_dict(include_input=is_sample))
+            ).to_dict(is_sample=is_sample))
         except Exception as e:
             results.append(TestResult(
                 expected=test['expected'],
                 passed=False,
+                logs=new_stdout.getvalue(),
                 error=traceback.format_exc(),
                 input=test['input'],
-            ).to_dict(include_input=is_sample))
+            ).to_dict(is_sample=is_sample))
+        finally:
+            sys.stdout = old_stdout
             
     return {{
         "test_results": results,
@@ -94,6 +107,7 @@ def run_tests(solution, method_name, test_data, is_sample: bool = False):
 import com.google.gson.*;
 import java.lang.reflect.*;
 import java.util.*;
+import java.io.*;
 
 {code}
 
@@ -115,11 +129,32 @@ def run_tests(solution, method_name, test_data, is_sample: bool = False):
         }}
     }}
 
+    private static class LogCapture {{
+        private final ByteArrayOutputStream baos;
+        private final PrintStream original;
+        private final PrintStream capture;
+
+        public LogCapture() {{
+            this.baos = new ByteArrayOutputStream();
+            this.original = System.out;
+            this.capture = new PrintStream(baos);
+        }}
+
+        public void start() {{
+            System.setOut(capture);
+        }}
+
+        public String stop() {{
+            System.setOut(original);
+            return baos.toString().trim();
+        }}
+    }}
+
     private static boolean compare(Object result, Object expected) {{
         {compare_func}
     }}
 
-    public static JsonArray runTests(Solution solution, JsonArray testData, boolean showInput) {{
+    public static JsonArray runTests(Solution solution, JsonArray testData, boolean isSample) {{
         JsonArray results = new JsonArray();
         Gson gson = new Gson();
 
@@ -146,21 +181,28 @@ def run_tests(solution, method_name, test_data, is_sample: bool = False):
         for (JsonElement testElement : testData) {{
             JsonObject test = testElement.getAsJsonObject();
             JsonObject result = new JsonObject();
+            LogCapture logCapture = new LogCapture();
 
             try {{
                 String inputStr = test.get("input").getAsString();
                 Object[] args = parseArguments(inputStr, paramTypes);
+
+                logCapture.start();
                 Object output = targetMethod.invoke(solution, args);
+                String logs = logCapture.stop();
+
                 Object expected = parseValue(test.get("expected").getAsString(), returnType);
 
                 boolean passed = compare(output, expected);
                 result.addProperty("passed", passed);
                 result.addProperty("output", gson.toJson(output));
                 result.addProperty("expected", test.get("expected").getAsString());
-                if (showInput) {{
+                if (isSample) {{
+                    result.addProperty("logs", logs);
                     result.addProperty("input", inputStr);
                 }}
             }} catch (Exception e) {{
+                logCapture.stop();
                 Throwable cause = e;
                 if (e instanceof InvocationTargetException && e.getCause() != null) {{
                     cause = e.getCause();
@@ -459,6 +501,7 @@ def run_tests(solution, method_name, test_data, is_sample: bool = False):
 #include <type_traits>
 #include <bits/stdc++.h>
 #include <vector>
+#include <streambuf>
 
 using namespace std;
 
@@ -585,14 +628,19 @@ def run_tests(solution, method_name, test_data, is_sample: bool = False):
     {compare_func}
 }}
 
-Json::Value runTests(Solution& solution, const Json::Value& testData, bool showInput) {{
+Json::Value runTests(Solution& solution, const Json::Value& testData, bool isSample) {{
     Json::Value results(Json::arrayValue);
     Json::CharReaderBuilder builder;
     Json::CharReader* reader = builder.newCharReader();
     string errors;
 
     for (const auto &test : testData) {{
         Json::Value testResult;
+
+        stringstream logStream;
+        streambuf* oldCout = cout.rdbuf(); 
+        cout.rdbuf(logStream.rdbuf()); 
+
         try {{
             auto args = parseArguments(test["input"].asString());
 
@@ -615,13 +663,15 @@ def run_tests(solution, method_name, test_data, is_sample: bool = False):
             writer["indentation"] = "";
             testResult["output"] = Json::writeString(writer, output_json);
             testResult["expected"] = Json::writeString(writer, expected);
-            if (showInput) {{
+            if (isSample) {{
+                testResult["logs"] = logStream.str(); 
                 testResult["input"] = test["input"];
             }}
         }} catch (const exception &e) {{
             testResult["error"] = e.what();
             testResult["passed"] = false;
         }}
+        cout.rdbuf(oldCout);
         results.append(testResult);
     }}
     delete reader;

app/services/execution/test_generator.py

@@ -30,12 +30,12 @@ def generate_test_file(
         """
 
     @abstractmethod
-    def get_file_extension(self, lang: str) -> str:
-        """
-        Get the file extension for the given language.
+    def get_file_extension(self) -> str:
+        """Get the file extension for the given language."""
 
-        :param lang: The language.
-        """
+    @abstractmethod
+    def get_line_offset(self) -> int:
+        """Get the line offset for the given language."""
 
     def process_quotes(self, json_str: str) -> str:
         """
@@ -67,9 +67,12 @@ def generate_test_file(
             sample_data=json.dumps(sample_data),
         )
 
-    def get_file_extension(self, lang: str) -> str:
+    def get_file_extension(self) -> str:
         return ".py"
 
+    def get_line_offset(self) -> int:
+        return 6
+
 
 class JavaTestGenerator(TestGenerator):
     def generate_test_file(
@@ -102,9 +105,6 @@ def generate_test_file(
             sample_data=sample_data,
         )
 
-    def get_file_extension(self, lang: str) -> str:
-        return ".java"
-
     def data_chunks(self, data: str) -> str:
         """
         Converts the data to a string of concatenated JSON strings.
@@ -125,6 +125,12 @@ def data_chunks(self, data: str) -> str:
         chunks = "".join(['.append("' + s + '")' for s in json_strings])
         return f"new StringBuilder(){chunks}.toString()"
 
+    def get_file_extension(self) -> str:
+        return ".java"
+
+    def get_line_offset(self) -> int:
+        return 8
+
 
 class CppTestGenerator(TestGenerator):
     def generate_test_file(
@@ -198,5 +204,8 @@ def process_test_data(self, test_data: List[Dict]) -> List[Dict]:
             processed_data.append(data)
         return processed_data
 
-    def get_file_extension(self, lang: str) -> str:
+    def get_file_extension(self) -> str:
         return ".cpp"
+
+    def get_line_offset(self) -> int:
+        return 14

app/services/execution/types.py

@@ -1,4 +1,36 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Any, Optional
+
+
+class TestResult:
+    def __init__(
+        self,
+        expected: str,
+        passed: bool,
+        output: Any = None,
+        logs: str = None,
+        error: str = None,
+        input: str = None,
+    ):
+        self.expected = expected
+        self.output = str(output) if output is not None else None
+        self.passed = passed
+        self.logs = logs
+        self.error = error
+        self.input = input
+
+    def to_dict(self, is_sample: bool = True):
+        result = {
+            {
+                "expected": self.expected,
+                "output": self.output,
+                "passed": self.passed,
+                "error": self.error,
+            }
+        }
+        if is_sample:
+            result["logs"] = self.logs
+            result["input"] = self.input
+        return result
 
 
 class ExecutionResult:
@@ -10,13 +42,15 @@ def __init__(
         self,
         success: bool,
         message: Optional[str] = None,
-        test_results: Optional[List[Dict]] = None,
-        sample_results: Optional[List[Dict]] = None,
+        line_offset: Optional[int] = None,  # for error logs from templates
+        test_results: Optional[List[TestResult]] = None,
+        sample_results: Optional[List[TestResult]] = None,
         summary: Optional[Dict] = None,
         runtime_analysis: Optional[str] = None,
     ):
         self.success = success
         self.message = message
+        self.line_offset = line_offset
         self.test_results = test_results
         self.sample_results = sample_results
         self.summary = summary
@@ -39,6 +73,7 @@ def to_dict(self) -> Dict:
         return {
             "success": self.success,
             "message": self.message,
+            "line_offset": self.line_offset,
             "test_results": self.test_results,
             "sample_results": self.sample_results,
             "summary": self.summary