chore: import upstream snapshot with attribution

tests/unit_tests/tools/code/test_improve_code.py

@@ -0,0 +1,52 @@
+import pytest
+from unittest.mock import Mock, MagicMock
+from superagi.tools.code.improve_code import ImproveCodeTool
+
+@pytest.fixture
+def mock_improve_code_tool():
+    improve_code_tool = ImproveCodeTool()
+    improve_code_tool.resource_manager = Mock()
+    improve_code_tool.llm = Mock()
+    return improve_code_tool
+
+def test_execute(mock_improve_code_tool):
+    mock_improve_code_tool.resource_manager.get_files.return_value = ['test1', 'test2']
+    mock_improve_code_tool.resource_manager.read_file.return_value = "test file content"
+    mock_improve_code_tool.llm.chat_completion.return_value = {
+        "response":
+            {
+                "choices":
+                    [
+                        {
+                            "message":
+                                {
+                                    "content": "```\nimproved code\n```"
+                                }
+                        }
+                    ]
+            }
+    }
+    mock_improve_code_tool.resource_manager.write_file.return_value = "file saved successfully"
+
+    assert mock_improve_code_tool._execute() == "All codes improved and saved successfully in: test1 test2"
+
+def test_execute_with_error(mock_improve_code_tool):
+    mock_improve_code_tool.resource_manager.get_files.return_value = ['test1']
+    mock_improve_code_tool.resource_manager.read_file.return_value = "test file content"
+    mock_improve_code_tool.llm.chat_completion.return_value = {
+        "response":
+            {
+                "choices":
+                    [
+                        {
+                            "message":
+                                {
+                                    "content": "```\nimproved code\n```"
+                                }
+                        }
+                    ]
+            }
+    }
+    mock_improve_code_tool.resource_manager.write_file.return_value = "Error: Could not save file"
+
+    assert mock_improve_code_tool._execute() == "Error: Could not save file"

tests/unit_tests/tools/code/test_write_code.py

@@ -0,0 +1,40 @@
+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")

tests/unit_tests/tools/code/test_write_spec.py

@@ -0,0 +1,32 @@
+from unittest.mock import Mock
+
+import pytest
+
+from superagi.tools.code.write_spec import WriteSpecTool
+from unittest.mock import MagicMock
+
+
+class MockBaseLlm:
+    def chat_completion(self, messages, max_tokens):
+        return {"content": "Generated specification"}
+
+    def get_model(self):
+        return "gpt-3.5-turbo"
+
+class TestWriteSpecTool:
+
+    @pytest.fixture
+    def tool(self):
+        tool = WriteSpecTool()
+        tool.llm = MockBaseLlm()
+        tool.resource_manager = Mock()
+        mock_session = MagicMock(name="session")
+        tool.toolkit_config.session = mock_session
+        return tool
+
+    def test_execute(self, tool):
+        tool.resource_manager.write_file = Mock()
+        tool.resource_manager.write_file.return_value = "File write successful"
+        response = tool._execute("Test task description", "test_spec_file.txt")
+        assert response == "Generated specification\nSpecification generated and saved successfully"
+        tool.resource_manager.write_file.assert_called_once_with("test_spec_file.txt", "Generated specification")

tests/unit_tests/tools/code/test_write_test.py

@@ -0,0 +1,53 @@
+from unittest.mock import Mock, patch
+
+from superagi.tools.code.write_test import WriteTestTool
+from unittest.mock import MagicMock
+
+
+def test_write_test_tool_init():
+    tool = WriteTestTool()
+    assert tool.llm is None
+    assert tool.agent_id is None
+    assert tool.name == "WriteTestTool"
+    assert tool.description is not None
+    assert tool.goals == []
+    assert tool.resource_manager is None
+
+@patch('superagi.tools.code.write_test.PromptReader')
+@patch('superagi.tools.code.write_test.AgentPromptBuilder')
+@patch('superagi.tools.code.write_test.TokenCounter')
+def test_execute(mock_token_counter, mock_agent_prompt_builder, mock_prompt_reader):
+    test_tool = WriteTestTool()
+    test_tool.tool_response_manager = Mock()
+    test_tool.resource_manager = Mock()
+    test_tool.llm = Mock()
+    mock_session = MagicMock(name="session")
+    test_tool.toolkit_config.session = mock_session
+
+    test_tool.tool_response_manager.get_last_response.return_value = 'WriteSpecTool response'
+    mock_prompt_reader.read_tools_prompt.return_value = 'Prompt template {goals} {test_description} {spec}'
+    mock_agent_prompt_builder.add_list_items_to_string.return_value = 'Goals string'
+    test_tool.llm.get_model.return_value = 'Model'
+    mock_token_counter.count_message_tokens.return_value = 100
+    mock_token_counter.token_limit.return_value = 1000
+    test_tool.llm.chat_completion.return_value = {
+        'content': 'File1\n```\nCode1```File2\n```\nCode2```',
+    }
+    test_tool.resource_manager.write_file.return_value = 'Success'
+
+    result = test_tool._execute('Test description', 'test_file.py')
+
+    assert 'File1' in result
+    assert 'Code1' in result
+    assert 'File2' in result
+    assert 'Code2' in result
+    assert 'Tests generated and saved successfully in test_file.py' in result
+
+    mock_prompt_reader.read_tools_prompt.assert_called_once()
+    mock_agent_prompt_builder.add_list_items_to_string.assert_called_once_with(test_tool.goals)
+    test_tool.tool_response_manager.get_last_response.assert_called()
+    test_tool.llm.get_model.assert_called()
+    mock_token_counter.count_message_tokens.assert_called()
+    mock_token_counter().token_limit.assert_called()
+    test_tool.llm.chat_completion.assert_called()
+    assert test_tool.resource_manager.write_file.call_count == 2