unit test clean up

unittesting/boss.py

@@ -1,72 +1,66 @@
-import pytest
+import unittest
-
+from unittest.mock import patch
-def test_WorkerNode_create_agent():
+
-    # assuming llm, tools, and vectorstore are initialized properly
+class TestSwarms(unittest.TestCase):
-    worker_node = WorkerNode(llm, tools, vectorstore)
+
-    worker_node.create_agent('test_agent', 'test_role', False, {})
+    def test_WorkerNode_create_agent(self):
-    assert worker_node.agent is not None
+        worker_node = WorkerNode(llm, tools, vectorstore)
-    assert worker_node.agent.chain.verbose
+        worker_node.create_agent('test_agent', 'test_role', False, {})
-
+        self.assertIsNotNone(worker_node.agent)
-def test_WorkerNode_run_agent():
+        self.assertEqual(worker_node.agent.ai_name, 'test_agent')
-    worker_node = WorkerNode(llm, tools, vectorstore)
+        self.assertEqual(worker_node.agent.ai_role, 'test_role')
-    worker_node.create_agent('test_agent', 'test_role', False, {})
+        self.assertEqual(worker_node.agent.human_in_the_loop, False)
-    worker_node.run_agent('test prompt')  # check it runs without error
+        self.assertTrue(worker_node.agent.chain.verbose)
-
+
-def test_BossNode_create_task():
+    def test_WorkerNode_tools(self):
-    # assuming llm, vectorstore, task_execution_chain are initialized properly
+        worker_node = WorkerNode(llm, tools, vectorstore)
-    boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
+        worker_node.create_agent('test_agent', 'test_role', False, {})
-    task = boss_node.create_task('test task')
+        for tool in worker_node.tools:
-    assert task == {'objective': 'test task'}
+            self.assertIsNotNone(tool)
-
+
-def test_BossNode_execute_task():
+    @patch.object(WorkerNode, 'run_agent')
-    boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
+    def test_WorkerNode_run_agent_called(self, mock_run_agent):
-    task = boss_node.create_task('test task')
+        worker_node = WorkerNode(llm, tools, vectorstore)
-    boss_node.execute_task(task)  # check it runs without error
+        worker_node.create_agent('test_agent', 'test_role', False, {})
-
+        worker_node.run_agent('test_prompt')
-
+        mock_run_agent.assert_called_once()
-def test_WorkerNode_tools():
+
-    worker_node = WorkerNode(llm, tools, vectorstore)
+    def test_BossNode_create_task(self):
-    worker_node.create_agent('test_agent', 'test_role', False, {})
+        boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
-    
+        task = boss_node.create_task('objective')
-    # Check that all tools are instantiated
+        self.assertEqual(task, {"objective": "objective"})
-    for tool in worker_node.tools:
+
-        assert tool is not None
+    def test_BossNode_AgentExecutor(self):
-
+        boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
-def test_BossNode_AgentExecutor():
+        self.assertIsNotNone(boss_node.baby_agi.task_execution_chain)
-    boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
+
-    
+    def test_BossNode_LLMChain(self):
-    # Check that the AgentExecutor is correctly initialized
+        boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
-    assert boss_node.baby_agi.task_execution_chain is not None
+        self.assertIsNotNone(boss_node.baby_agi.task_execution_chain.agent.llm_chain)
-
+
-def test_BossNode_LLMChain():
+    @patch.object(BossNode, 'execute_task')
-    boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
+    def test_BossNode_execute_task_called_with_correct_arg(self, mock_execute_task):
-    
+        boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
-    # Check that the LLMChain in ZeroShotAgent is working
+        task = boss_node.create_task('objective')
-    assert boss_node.baby_agi.task_execution_chain.agent.llm_chain is not None
+        boss_node.execute_task(task)
-
+        mock_execute_task.assert_called_once_with(task)
-
+
-
+    @patch.object(BossNode, 'execute_task')
-def test_WorkerNode_create_agent():
+    def test_BossNode_execute_task_output(self, mock_execute_task):
-    worker_node = WorkerNode(llm, tools, vectorstore)
+        mock_execute_task.return_value = "some_output"
-    worker_node.create_agent('test_agent', 'test_role', False, {})
+        boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
-
+        task = boss_node.create_task('objective')
-    assert worker_node.agent.ai_name == 'test_agent'
+        output = boss_node.execute_task(task)
-    assert worker_node.agent.ai_role == 'test_role'
+        self.assertIsNotNone(output)
-    assert worker_node.agent.human_in_the_loop == False
+
-
+    def test_BossNode_execute_task_error_handling(self):
-def test_BossNode_execute_task_output():
+        boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
-    boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
+        try:
-    task = boss_node.create_task('objective')
+            boss_node.execute_task(None)
-
+            self.assertTrue(True)
-    # The output of the execute_task method should not be None
+        except Exception:
-    assert boss_node.execute_task(task) is not None
+            self.fail("boss_node.execute_task raised Exception unexpectedly!")
-
+
-def test_BossNode_execute_task_error_handling():
+# Run the tests
-    boss_node = BossNode(llm, vectorstore, task_execution_chain, False, 3)
+if __name__ == '__main__':
-
+    unittest.main()
-    # Try executing an invalid task and make sure it doesn't crash
-    try:
-        boss_node.execute_task(None)
-        assert True
-    except Exception:
-        assert False