ai-mirow-projects
/
swarms
mirror of https://github.com/kyegomez/swarms
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

CLEANUP: tests/workers

Browse Source
pull/176/head
Kye 1 year ago
parent 5354c97341
commit cfd2b2e7cd
5 changed files with 0 additions and 328 deletions
Show Stats Download Patch File Download Diff File

34
tests/workers/multi_model_worker.py
Unescape View File

@ -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)

58
tests/workers/omni_worker.py
Unescape View File

@ -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)

51
tests/workers/worker_agent_ultra.py
Unescape View File

@ -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)

94
tests/workers/worker_node.py
Unescape View File

@ -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()

91
tests/workers/worker_ultra.py
Unescape View File

@ -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
Write Preview
Loading…
Cancel
Save