diff --git a/tests/agents/workers/worker_ultra.py b/tests/agents/workers/worker_ultra.py index e69de29b..b9fb523d 100644 --- a/tests/agents/workers/worker_ultra.py +++ b/tests/agents/workers/worker_ultra.py @@ -0,0 +1,63 @@ +import pytest +from unittest.mock import Mock, patch +from swarms.workers.worker_agent_ultra import WorkerUltraNode, WorkerUltraNodeInitializer # replace 'yourmodule' with the actual module name + +@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') as mock_llm, patch.object(worker_ultra_node_initializer, 'initialize_toolsets') as mock_toolsets, patch.object(worker_ultra_node_initializer, 'initialize_vectorstore') as mock_vectorstore: + worker_node = worker_ultra_node_initializer.create_worker_node() + assert worker_node is not None