|
|
|
@ -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
|
|
|
|
|