From 6cc27795cd2d585271f7ca0db69da235faddba8e Mon Sep 17 00:00:00 2001 From: Kye Date: Fri, 30 Jun 2023 18:33:35 -0400 Subject: [PATCH] test swarms clean up --- unittesting/swarms.py | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/unittesting/swarms.py b/unittesting/swarms.py index 89bccc97..aa8488fa 100644 --- a/unittesting/swarms.py +++ b/unittesting/swarms.py @@ -1,11 +1,13 @@ import unittest from unittest.mock import patch -from swarms.agents.swarms import llm, tools, vectorstore, task_execution_chain, WorkerNode, BossNode +# from swarms.agents.swarms import llm, tools, vectorstore, task_execution_chain, WorkerNode, BossNode +from swarms.agents.swarms import worker_node, boss_node + class TestSwarms(unittest.TestCase): def test_WorkerNode_create_agent(self): - worker_node = WorkerNode(llm, tools, vectorstore) + # worker_node = WorkerNode(llm, tools, vectorstore) worker_node.create_agent('test_agent', 'test_role', False, {}) self.assertIsNotNone(worker_node.agent) self.assertEqual(worker_node.agent.ai_name, 'test_agent') @@ -14,48 +16,40 @@ class TestSwarms(unittest.TestCase): self.assertTrue(worker_node.agent.chain.verbose) def test_WorkerNode_tools(self): - worker_node = WorkerNode(llm, tools, vectorstore) worker_node.create_agent('test_agent', 'test_role', False, {}) for tool in worker_node.tools: self.assertIsNotNone(tool) - @patch.object(WorkerNode, 'run_agent') + @patch.object(worker_node, 'run_agent') def test_WorkerNode_run_agent_called(self, mock_run_agent): - worker_node = WorkerNode(llm, tools, vectorstore) worker_node.create_agent('test_agent', 'test_role', False, {}) worker_node.run_agent('test_prompt') mock_run_agent.assert_called_once() def test_BossNode_create_task(self): - boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3) task = boss_node.create_task('objective') self.assertEqual(task, {"objective": "objective"}) def test_BossNode_AgentExecutor(self): - boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3) self.assertIsNotNone(boss_node.baby_agi.task_execution_chain) def test_BossNode_LLMChain(self): - boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3) self.assertIsNotNone(boss_node.baby_agi.task_execution_chain.agent.llm_chain) - @patch.object(BossNode, 'execute_task') + @patch.object(boss_node, 'execute_task') def test_BossNode_execute_task_called_with_correct_arg(self, mock_execute_task): - boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3) task = boss_node.create_task('objective') boss_node.execute_task(task) mock_execute_task.assert_called_once_with(task) - @patch.object(BossNode, 'execute_task') + @patch.object(boss_node, 'execute_task') def test_BossNode_execute_task_output(self, mock_execute_task): mock_execute_task.return_value = "some_output" - boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3) task = boss_node.create_task('objective') output = boss_node.execute_task(task) self.assertIsNotNone(output) def test_BossNode_execute_task_error_handling(self): - boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3) try: boss_node.execute_task(None) self.assertTrue(True)