wehub-resource-sync
40 lines
1.7 KiB
Python
40 lines
1.7 KiB
Python
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from superagi.resource_manager.file_manager import FileManager
|
|
from superagi.tools.code.write_code import CodingTool
|
|
from superagi.tools.tool_response_query_manager import ToolResponseQueryManager
|
|
from unittest.mock import MagicMock
|
|
|
|
|
|
class MockBaseLlm:
|
|
def chat_completion(self, messages, max_tokens):
|
|
return {"content": "File1.py\n```python\nprint('Hello World')\n```\n\nFile2.py\n```python\nprint('Hello again')\n```"}
|
|
|
|
def get_model(self):
|
|
return "gpt-3.5-turbo"
|
|
|
|
class TestCodingTool:
|
|
|
|
@pytest.fixture
|
|
def tool(self):
|
|
tool = CodingTool()
|
|
tool.llm = MockBaseLlm()
|
|
tool.resource_manager = Mock(spec=FileManager)
|
|
tool.tool_response_manager = Mock(spec=ToolResponseQueryManager)
|
|
mock_session = MagicMock(name="session")
|
|
tool.toolkit_config.session = mock_session
|
|
return tool
|
|
|
|
def test_execute(self, tool):
|
|
tool.resource_manager.write_file.return_value = "File write successful"
|
|
tool.tool_response_manager.get_last_response.return_value = "Mocked Spec"
|
|
|
|
response = tool._execute("Test spec description")
|
|
assert response == "File1.py\n```python\nprint('Hello World')\n```\n\nFile2.py\n```python\nprint('Hello again')\n```\n Codes generated and saved successfully in File1.py, File2.py"
|
|
|
|
tool.resource_manager.write_file.assert_any_call("README.md", 'File1.py\n')
|
|
tool.resource_manager.write_file.assert_any_call("File1.py", "print('Hello World')\n")
|
|
tool.resource_manager.write_file.assert_any_call("File2.py", "print('Hello again')\n")
|
|
tool.tool_response_manager.get_last_response.assert_called_once_with("WriteSpecTool") |