fix: enhance output handling and add comprehensive tests for code execution tool

src/python_tool.py

@@ -172,7 +172,10 @@ class CodeExecutionTool:
             
             output = output_buffer.getvalue()
             if len(output) > self.max_output_size:
-                output = output[:self.max_output_size] + "... [output truncated]"
+                truncation_message = f"\n... [output truncated to {self.max_output_size} characters]"
+                output = output[:self.max_output_size - len(truncation_message)] + truncation_message
+            else:
+                output = output.strip()
             
             # Extract the numeric value
             numeric_result = self._extract_numeric_value(output)

tests/test_python_tool.py

@@ -40,5 +40,59 @@ class TestCodeExecutionTool(unittest.TestCase):
         self.assertEqual(reversed_question, "If you understand this sentence, write the opposite of the word \"left\" as the answer.")
         self.assertEqual(expected_answer, "Right")
 
+    def test_execute_code_success(self):
+        """Test successful execution of safe Python code."""
+        safe_code = "print(42)"
+        result = self.code_tool.execute_code(safe_code)
+        self.assertTrue(result["success"])
+        self.assertEqual(result["raw_output"].strip(), "42")
+
+    def test_execute_code_numeric_extraction(self):
+        """Test numeric value extraction from code output."""
+        numeric_code = "print(3.14)"
+        result = self.code_tool.execute_code(numeric_code)
+        self.assertTrue(result["success"])
+        self.assertTrue(result["has_numeric_result"])
+        self.assertEqual(result["numeric_value"], 3.14)
+
+    def test_execute_code_timeout(self):
+        """Test that code execution times out as expected."""
+        timeout_code = "while True: pass"
+        result = self.code_tool.execute_code(timeout_code)
+        self.assertFalse(result["success"])
+        self.assertIn("timed out", result["error"].lower())
+
+    def test_execute_code_error_handling(self):
+        """Test error handling for code that raises exceptions."""
+        error_code = "raise ValueError('Test error')"
+        result = self.code_tool.execute_code(error_code)
+        self.assertFalse(result["success"])
+        self.assertIn("ValueError", result["error"])
+
+    def test_execute_code_output_size_limit(self):
+        """Test that output is truncated if it exceeds max_output_size."""
+        large_output_code = "print('A' * 20000)"
+        result = self.code_tool.execute_code(large_output_code)
+        self.assertTrue(result["success"])
+        self.assertLessEqual(len(result["raw_output"]), self.code_tool.max_output_size)
+
+    def test_execute_file_success(self):
+        """Test successful execution of a Python file."""
+        test_file_path = "test_script.py"
+        with open(test_file_path, "w") as f:
+            f.write("print('File executed successfully')")
+        
+        result = self.code_tool.execute_file(test_file_path)
+        self.assertTrue(result["success"])
+        self.assertEqual(result["raw_output"].strip(), "File executed successfully")
+        
+        os.remove(test_file_path)
+
+    def test_execute_file_not_found(self):
+        """Test handling of file not found error."""
+        result = self.code_tool.execute_file("non_existent_file.py")
+        self.assertFalse(result["success"])
+        self.assertIn("File not found", result["error"])
+
 if __name__ == "__main__":
     unittest.main()