removing agent prompt and adding unit tests for task queue

jroell · Jun 5, 2023 · bea0062 · bea0062
1 parent c6a5fc0
commit bea0062
Show file tree

Hide file tree

Showing 6 changed files with 54 additions and 149 deletions.
diff --git a/gui/package-lock.json b/gui/package-lock.json
diff --git a/superagi/agent/agent_prompt.py b/superagi/agent/agent_prompt.py
diff --git a/superagi/agent/agent_prompt_builder.py b/superagi/agent/agent_prompt_builder.py
@@ -1,10 +1,10 @@
 import json
+import re
+
 from pydantic.types import List
-from superagi.agent.agent_prompt import AgentPrompt
+
 from superagi.helper.token_counter import TokenCounter
 from superagi.tools.base_tool import BaseTool
-from fastapi_sqlalchemy import db
-import re
 
 FINISH_NAME = "finish"
 

diff --git a/superagi/agent/task_listener.py b/superagi/agent/task_listener.py
diff --git a/tests/agent/__init__.py b/tests/agent/__init__.py
diff --git a/tests/agent/test_task_queue.py b/tests/agent/test_task_queue.py
@@ -0,0 +1,51 @@
+import unittest
+from unittest.mock import patch
+
+from superagi.agent.task_queue import TaskQueue
+
+
+class TaskQueueTests(unittest.TestCase):
+    def setUp(self):
+        self.queue_name = "test_queue"
+        self.queue = TaskQueue(self.queue_name)
+
+    @patch.object(TaskQueue, 'add_task')
+    def test_add_task(self, mock_add_task):
+        task = "Do something"
+        self.queue.add_task(task)
+        mock_add_task.assert_called_with(task)
+
+    @patch.object(TaskQueue, 'complete_task')
+    def test_complete_task(self, mock_complete_task):
+        task = "Do something"
+        response = "Task completed"
+        self.queue.complete_task(response)
+        mock_complete_task.assert_called_with(response)
+
+    @patch.object(TaskQueue, 'get_first_task')
+    def test_get_first_task(self, mock_get_first_task):
+        self.queue.get_first_task()
+        mock_get_first_task.assert_called()
+
+    @patch.object(TaskQueue, 'get_tasks')
+    def test_get_tasks(self, mock_get_tasks):
+        self.queue.get_tasks()
+        mock_get_tasks.assert_called()
+
+    @patch.object(TaskQueue, 'get_completed_tasks')
+    def test_get_completed_tasks(self, mock_get_completed_tasks):
+        self.queue.get_completed_tasks()
+        mock_get_completed_tasks.assert_called()
+
+    @patch.object(TaskQueue, 'clear_tasks')
+    def test_clear_tasks(self, mock_clear_tasks):
+        self.queue.clear_tasks()
+        mock_clear_tasks.assert_called()
+
+    @patch.object(TaskQueue, 'get_last_task_details')
+    def test_get_last_task_details(self, mock_get_last_task_details):
+        self.queue.get_last_task_details()
+        mock_get_last_task_details.assert_called()
+
+if __name__ == '__main__':
+    unittest.main()