From cfd2b2e7cd53f3e6e8aff1c4dc08336a1f2159f0 Mon Sep 17 00:00:00 2001 From: Kye Date: Thu, 23 Nov 2023 15:21:00 -0800 Subject: [PATCH] CLEANUP: tests/workers --- tests/workers/multi_model_worker.py | 34 ----------- tests/workers/omni_worker.py | 58 ------------------ tests/workers/worker_agent_ultra.py | 51 ---------------- tests/workers/worker_node.py | 94 ----------------------------- tests/workers/worker_ultra.py | 91 ---------------------------- 5 files changed, 328 deletions(-) delete mode 100644 tests/workers/multi_model_worker.py delete mode 100644 tests/workers/omni_worker.py delete mode 100644 tests/workers/worker_agent_ultra.py delete mode 100644 tests/workers/worker_node.py delete mode 100644 tests/workers/worker_ultra.py diff --git a/tests/workers/multi_model_worker.py b/tests/workers/multi_model_worker.py deleted file mode 100644 index f011d642..00000000 --- a/tests/workers/multi_model_worker.py +++ /dev/null @@ -1,34 +0,0 @@ -import pytest -from unittest.mock import Mock -from swarms.agents.multi_modal_agent import ( - MultiModalVisualAgent, - MultiModalVisualAgentTool, -) - - -@pytest.fixture -def multimodal_agent(): - # Mock the MultiModalVisualAgent - mock_agent = Mock(spec=MultiModalVisualAgent) - mock_agent.run_text.return_value = "Expected output from agent" - return mock_agent - - -@pytest.fixture -def multimodal_agent_tool(multimodal_agent): - # Use the mocked MultiModalVisualAgent in the MultiModalVisualAgentTool - return MultiModalVisualAgentTool(multimodal_agent) - - -@pytest.mark.parametrize( - "text_input, expected_output", - [ - ("Hello, world!", "Expected output from agent"), - ("Another task", "Expected output from agent"), - ], -) -def test_run(multimodal_agent_tool, text_input, expected_output): - assert multimodal_agent_tool._run(text_input) == expected_output - - # You can also test if the MultiModalVisualAgent's run_text method was called with the right argument - multimodal_agent_tool.agent.run_text.assert_called_with(text_input) diff --git a/tests/workers/omni_worker.py b/tests/workers/omni_worker.py deleted file mode 100644 index 0557285d..00000000 --- a/tests/workers/omni_worker.py +++ /dev/null @@ -1,58 +0,0 @@ -import pytest - -from swarms.worker.omni_worker import OmniWorkerAgent - - -@pytest.fixture -def omni_worker(): - api_key = "test-key" - api_endpoint = "test-endpoint" - api_type = "test-type" - return OmniWorkerAgent(api_key, api_endpoint, api_type) - - -@pytest.mark.parametrize( - "data, expected_response", - [ - ( - { - "messages": ["Hello"], - "api_key": "key1", - "api_type": "type1", - "api_endpoint": "endpoint1", - }, - {"response": "Hello back from Huggingface!"}, - ), - ( - { - "messages": ["Goodbye"], - "api_key": "key2", - "api_type": "type2", - "api_endpoint": "endpoint2", - }, - {"response": "Goodbye from Huggingface!"}, - ), - ], -) -def test_chat_valid_data(mocker, omni_worker, data, expected_response): - mocker.patch( - "yourmodule.chat_huggingface", return_value=expected_response - ) # replace 'yourmodule' with actual module name - assert omni_worker.chat(data) == expected_response - - -@pytest.mark.parametrize( - "invalid_data", - [ - {"messages": ["Hello"]}, # missing api_key, api_type and api_endpoint - {"messages": ["Hello"], "api_key": "key1"}, # missing api_type and api_endpoint - { - "messages": ["Hello"], - "api_key": "key1", - "api_type": "type1", - }, # missing api_endpoint - ], -) -def test_chat_invalid_data(omni_worker, invalid_data): - with pytest.raises(ValueError): - omni_worker.chat(invalid_data) diff --git a/tests/workers/worker_agent_ultra.py b/tests/workers/worker_agent_ultra.py deleted file mode 100644 index 3cf112a2..00000000 --- a/tests/workers/worker_agent_ultra.py +++ /dev/null @@ -1,51 +0,0 @@ -import pytest -from unittest.mock import Mock -from swarms.workers.worker_agent_ultra import WorkerUltraNode # import your module here - - -def test_create_agent(): - mock_llm = Mock() - mock_toolset = {"test_toolset": Mock()} - mock_vectorstore = Mock() - worker = WorkerUltraNode(mock_llm, mock_toolset, mock_vectorstore) - worker.create_agent() - assert worker.agent is not None - - -@pytest.mark.parametrize("invalid_toolset", [123, "string", 0.45]) -def test_add_toolset_invalid(invalid_toolset): - mock_llm = Mock() - mock_toolset = {"test_toolset": Mock()} - mock_vectorstore = Mock() - worker = WorkerUltraNode(mock_llm, mock_toolset, mock_vectorstore) - with pytest.raises(TypeError): - worker.add_toolset(invalid_toolset) - - -@pytest.mark.parametrize("invalid_prompt", [123, None, "", []]) -def test_run_invalid_prompt(invalid_prompt): - mock_llm = Mock() - mock_toolset = {"test_toolset": Mock()} - mock_vectorstore = Mock() - worker = WorkerUltraNode(mock_llm, mock_toolset, mock_vectorstore) - with pytest.raises((TypeError, ValueError)): - worker.run(invalid_prompt) - - -def test_run_valid_prompt(mocker): - mock_llm = Mock() - mock_toolset = {"test_toolset": Mock()} - mock_vectorstore = Mock() - worker = WorkerUltraNode(mock_llm, mock_toolset, mock_vectorstore) - mocker.patch.object(worker, "create_agent") - assert worker.run("Test prompt") == "Task completed by WorkerNode" - - -def test_worker_node(): - worker = worker_ultra_node("test-key") - assert isinstance(worker, WorkerUltraNode) - - -def test_worker_node_no_key(): - with pytest.raises(ValueError): - worker_ultra_node(None) diff --git a/tests/workers/worker_node.py b/tests/workers/worker_node.py deleted file mode 100644 index e97b5023..00000000 --- a/tests/workers/worker_node.py +++ /dev/null @@ -1,94 +0,0 @@ -import pytest -from unittest.mock import MagicMock, patch -from swarms.worker.worker_node import ( - WorkerNodeInitializer, - WorkerNode, -) # replace your_module with actual module name - - -# Mock Tool for testing -class MockTool(Tool): - pass - - -# Fixture for llm -@pytest.fixture -def mock_llm(): - return MagicMock() - - -# Fixture for vectorstore -@pytest.fixture -def mock_vectorstore(): - return MagicMock() - - -# Fixture for Tools -@pytest.fixture -def mock_tools(): - return [MockTool(), MockTool(), MockTool()] - - -# Fixture for WorkerNodeInitializer -@pytest.fixture -def worker_node(mock_llm, mock_tools, mock_vectorstore): - return WorkerNodeInitializer( - llm=mock_llm, tools=mock_tools, vectorstore=mock_vectorstore - ) - - -# Fixture for WorkerNode -@pytest.fixture -def mock_worker_node(): - return WorkerNode(openai_api_key="test_api_key") - - -# WorkerNodeInitializer Tests -def test_worker_node_init(worker_node): - assert worker_node.llm is not None - assert worker_node.tools is not None - assert worker_node.vectorstore is not None - - -def test_worker_node_create_agent(worker_node): - with patch.object(AutoGPT, "from_llm_and_tools") as mock_method: - worker_node.create_agent() - mock_method.assert_called_once() - - -def test_worker_node_add_tool(worker_node): - initial_tools_count = len(worker_node.tools) - new_tool = MockTool() - worker_node.add_tool(new_tool) - assert len(worker_node.tools) == initial_tools_count + 1 - - -def test_worker_node_run(worker_node): - with patch.object(worker_node.agent, "run") as mock_run: - worker_node.run(prompt="test prompt") - mock_run.assert_called_once() - - -# WorkerNode Tests -def test_worker_node_llm(mock_worker_node): - with patch.object(mock_worker_node, "initialize_llm") as mock_method: - mock_worker_node.initialize_llm(llm_class=MagicMock(), temperature=0.5) - mock_method.assert_called_once() - - -def test_worker_node_tools(mock_worker_node): - with patch.object(mock_worker_node, "initialize_tools") as mock_method: - mock_worker_node.initialize_tools(llm_class=MagicMock()) - mock_method.assert_called_once() - - -def test_worker_node_vectorstore(mock_worker_node): - with patch.object(mock_worker_node, "initialize_vectorstore") as mock_method: - mock_worker_node.initialize_vectorstore() - mock_method.assert_called_once() - - -def test_worker_node_create_worker_node(mock_worker_node): - with patch.object(mock_worker_node, "create_worker_node") as mock_method: - mock_worker_node.create_worker_node() - mock_method.assert_called_once() diff --git a/tests/workers/worker_ultra.py b/tests/workers/worker_ultra.py deleted file mode 100644 index b1485a28..00000000 --- a/tests/workers/worker_ultra.py +++ /dev/null @@ -1,91 +0,0 @@ -import pytest -from unittest.mock import Mock, patch -from swarms.workers.worker_agent_ultra import ( - WorkerUltraNode, - WorkerUltraNodeInitializer, -) - - -@pytest.fixture -def llm_mock(): - return Mock() - - -@pytest.fixture -def toolsets_mock(): - return Mock() - - -@pytest.fixture -def vectorstore_mock(): - return Mock() - - -@pytest.fixture -def worker_ultra_node(llm_mock, toolsets_mock, vectorstore_mock): - return WorkerUltraNode(llm_mock, toolsets_mock, vectorstore_mock) - - -def test_worker_ultra_node_create_agent(worker_ultra_node): - with patch("yourmodule.AutoGPT.from_llm_and_tools") as mock_method: - worker_ultra_node.create_agent() - mock_method.assert_called_once() - - -def test_worker_ultra_node_add_toolset(worker_ultra_node): - with pytest.raises(TypeError): - worker_ultra_node.add_toolset("wrong_toolset") - - -def test_worker_ultra_node_run(worker_ultra_node): - with patch.object(worker_ultra_node, "agent") as mock_agent: - mock_agent.run.return_value = None - result = worker_ultra_node.run("some prompt") - assert result == "Task completed by WorkerNode" - mock_agent.run.assert_called_once() - - -def test_worker_ultra_node_run_no_prompt(worker_ultra_node): - with pytest.raises(ValueError): - worker_ultra_node.run("") - - -@pytest.fixture -def worker_ultra_node_initializer(): - return WorkerUltraNodeInitializer("openai_api_key") - - -def test_worker_ultra_node_initializer_initialize_llm(worker_ultra_node_initializer): - with patch("yourmodule.ChatOpenAI") as mock_llm: - worker_ultra_node_initializer.initialize_llm(mock_llm) - mock_llm.assert_called_once() - - -def test_worker_ultra_node_initializer_initialize_toolsets( - worker_ultra_node_initializer, -): - with patch("yourmodule.Terminal"), patch("yourmodule.CodeEditor"), patch( - "yourmodule.RequestsGet" - ), patch("yourmodule.ExitConversation"): - toolsets = worker_ultra_node_initializer.initialize_toolsets() - assert len(toolsets) == 4 - - -def test_worker_ultra_node_initializer_initialize_vectorstore( - worker_ultra_node_initializer, -): - with patch("yourmodule.OpenAIEmbeddings"), patch( - "yourmodule.fauss.IndexFlatL2" - ), patch("yourmodule.FAISS"), patch("yourmodule.InMemoryDocstore"): - vectorstore = worker_ultra_node_initializer.initialize_vectorstore() - assert vectorstore is not None - - -def test_worker_ultra_node_initializer_create_worker_node( - worker_ultra_node_initializer, -): - with patch.object(worker_ultra_node_initializer, "initialize_llm"), patch.object( - worker_ultra_node_initializer, "initialize_toolsets" - ), patch.object(worker_ultra_node_initializer, "initialize_vectorstore"): - worker_node = worker_ultra_node_initializer.create_worker_node() - assert worker_node is not None