clean up

Former-commit-id: ccf47495b7

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