mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
54 lines
2.1 KiB
Python
54 lines
2.1 KiB
Python
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from crewai_tools.tools.code_interpreter_tool.code_interpreter_tool import (
|
|
CodeInterpreterTool,
|
|
)
|
|
|
|
|
|
class TestCodeInterpreterTool(unittest.TestCase):
|
|
@patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env")
|
|
def test_run_code_in_docker(self, docker_mock):
|
|
tool = CodeInterpreterTool()
|
|
code = "print('Hello, World!')"
|
|
libraries_used = ["numpy", "pandas"]
|
|
expected_output = "Hello, World!\n"
|
|
|
|
docker_mock().containers.run().exec_run().exit_code = 0
|
|
docker_mock().containers.run().exec_run().output = (
|
|
expected_output.encode()
|
|
)
|
|
result = tool.run_code_in_docker(code, libraries_used)
|
|
|
|
self.assertEqual(result, expected_output)
|
|
|
|
@patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env")
|
|
def test_run_code_in_docker_with_error(self, docker_mock):
|
|
tool = CodeInterpreterTool()
|
|
code = "print(1/0)"
|
|
libraries_used = ["numpy", "pandas"]
|
|
expected_output = "Something went wrong while running the code: \nZeroDivisionError: division by zero\n"
|
|
|
|
docker_mock().containers.run().exec_run().exit_code = 1
|
|
docker_mock().containers.run().exec_run().output = (
|
|
b"ZeroDivisionError: division by zero\n"
|
|
)
|
|
result = tool.run_code_in_docker(code, libraries_used)
|
|
|
|
self.assertEqual(result, expected_output)
|
|
|
|
@patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env")
|
|
def test_run_code_in_docker_with_script(self, docker_mock):
|
|
tool = CodeInterpreterTool()
|
|
code = """print("This is line 1")
|
|
print("This is line 2")"""
|
|
libraries_used = [] # No additional libraries needed for this test
|
|
expected_output = "This is line 1\nThis is line 2\n"
|
|
|
|
# Mock Docker responses
|
|
docker_mock().containers.run().exec_run().exit_code = 0
|
|
docker_mock().containers.run().exec_run().output = expected_output.encode()
|
|
|
|
result = tool.run_code_in_docker(code, libraries_used)
|
|
self.assertEqual(result, expected_output)
|