From 12970e59fbd4978d05b4e649a895625c5dc4b586 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Wed, 15 Nov 2023 17:19:37 -0500
Subject: [PATCH] autoscaler tests fix

---
 example.py                 |  9 ---------
 tests/swarms/autoscaler.py | 28 +++++++++++++++++++---------
 2 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/example.py b/example.py
index 5158d8d06..46e8b33ce 100644
--- a/example.py
+++ b/example.py
@@ -1,4 +1,3 @@
-from swarms.tools.tool import tool
 from swarms.models import OpenAIChat
 from swarms.structs import Flow
 
@@ -11,14 +10,6 @@
 )
 
 
-@tool
-def search_api(query: str):
-    """
-    This is a search API, you can search the web for information.
-    """
-    pass
-
-
 ## Initialize the workflow
 flow = Flow(
     llm=llm,
diff --git a/tests/swarms/autoscaler.py b/tests/swarms/autoscaler.py
index 951d3be75..bd884d73b 100644
--- a/tests/swarms/autoscaler.py
+++ b/tests/swarms/autoscaler.py
@@ -1,5 +1,15 @@
 from unittest.mock import patch
-from swarms.swarms.autoscaler import AutoScaler, Worker
+from swarms.swarms.autoscaler import AutoScaler
+from swarms.models import OpenAIChat
+from swarms.structs import Flow
+
+llm = OpenAIChat()
+
+flow = Flow(
+    llm=llm,
+    max_loops=2,
+    dashboard=True,
+)
 
 
 def test_autoscaler_initialization():
@@ -8,7 +18,7 @@ def test_autoscaler_initialization():
         scale_up_factor=2,
         idle_threshold=0.1,
         busy_threshold=0.8,
-        agent=Worker,
+        agent=flow,
     )
     assert isinstance(autoscaler, AutoScaler)
     assert autoscaler.scale_up_factor == 2
@@ -18,19 +28,19 @@ def test_autoscaler_initialization():
 
 
 def test_autoscaler_add_task():
-    autoscaler = AutoScaler(agent=Worker)
+    autoscaler = AutoScaler(agent=flow)
     autoscaler.add_task("task1")
     assert autoscaler.task_queue.qsize() == 1
 
 
 def test_autoscaler_scale_up():
-    autoscaler = AutoScaler(initial_agents=5, scale_up_factor=2, agent=Worker)
+    autoscaler = AutoScaler(initial_agents=5, scale_up_factor=2, agent=flow)
     autoscaler.scale_up()
     assert len(autoscaler.agents_pool) == 10
 
 
 def test_autoscaler_scale_down():
-    autoscaler = AutoScaler(initial_agents=5, agent=Worker)
+    autoscaler = AutoScaler(initial_agents=5, agent=flow)
     autoscaler.scale_down()
     assert len(autoscaler.agents_pool) == 4
 
@@ -38,7 +48,7 @@ def test_autoscaler_scale_down():
 @patch("your_module.AutoScaler.scale_up")
 @patch("your_module.AutoScaler.scale_down")
 def test_autoscaler_monitor_and_scale(mock_scale_down, mock_scale_up):
-    autoscaler = AutoScaler(initial_agents=5, agent=Worker)
+    autoscaler = AutoScaler(initial_agents=5, agent=flow)
     autoscaler.add_task("task1")
     autoscaler.monitor_and_scale()
     mock_scale_up.assert_called_once()
@@ -46,9 +56,9 @@ def test_autoscaler_monitor_and_scale(mock_scale_down, mock_scale_up):
 
 
 @patch("your_module.AutoScaler.monitor_and_scale")
-@patch("your_module.Worker.run")
+@patch("your_module.flow.run")
 def test_autoscaler_start(mock_run, mock_monitor_and_scale):
-    autoscaler = AutoScaler(initial_agents=5, agent=Worker)
+    autoscaler = AutoScaler(initial_agents=5, agent=flow)
     autoscaler.add_task("task1")
     autoscaler.start()
     mock_run.assert_called_once()
@@ -56,6 +66,6 @@ def test_autoscaler_start(mock_run, mock_monitor_and_scale):
 
 
 def test_autoscaler_del_agent():
-    autoscaler = AutoScaler(initial_agents=5, agent=Worker)
+    autoscaler = AutoScaler(initial_agents=5, agent=flow)
     autoscaler.del_agent()
     assert len(autoscaler.agents_pool) == 4