diff --git a/tests/agents/workers/worker_agent_ultra.py b/tests/agents/workers/worker_agent_ultra.py new file mode 100644 index 00000000..16bd3f5b --- /dev/null +++ b/tests/agents/workers/worker_agent_ultra.py @@ -0,0 +1,45 @@ +import pytest +from unittest.mock import Mock, MagicMock +from swarms.workers.worker_agent_ultra import WorkerUltraNode, WorkerUltraNodeInitializer # 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 = worker_ultra_node(None) diff --git a/tests/agents/workers/worker_node.py b/tests/agents/workers/worker_node.py index 2ae368e2..96340c6c 100644 --- a/tests/agents/workers/worker_node.py +++ b/tests/agents/workers/worker_node.py @@ -2,7 +2,7 @@ import pytest from swarms.agents.tools.agent_tools import * from unittest.mock import MagicMock, patch from pydantic import ValidationError -from swarms.worker.worker_nodegit add import WorkerNodeInitializer, WorkerNode # replace your_module with actual module name +from swarms.worker.worker_node import WorkerNodeInitializer, WorkerNode # replace your_module with actual module name # Mock Tool for testing diff --git a/tests/agents/workers/worker_ultra.py b/tests/agents/workers/worker_ultra.py new file mode 100644 index 00000000..e69de29b