parent
61c1894576
commit
8791248eff
@ -1,13 +0,0 @@
|
||||
from swarms.structs.autoscaler import AutoScaler
|
||||
from swarms.swarms.model_parallizer import ModelParallelizer
|
||||
from swarms.swarms.multi_agent_collab import MultiAgentCollaboration
|
||||
from swarms.swarms.base import AbstractSwarm
|
||||
|
||||
# from swarms.swarms.team import Team
|
||||
|
||||
__all__ = [
|
||||
"AutoScaler",
|
||||
"ModelParallelizer",
|
||||
"MultiAgentCollaboration",
|
||||
"AbstractSwarm",
|
||||
]
|
@ -1,5 +1,5 @@
|
||||
import pytest
|
||||
from swarms.swarms.model_parallizer import ModelParallelizer
|
||||
from swarms.structs.model_parallizer import ModelParallelizer
|
||||
from swarms.models import (
|
||||
HuggingfaceLLM,
|
||||
Mixtral,
|
@ -0,0 +1,52 @@
|
||||
import json
|
||||
import unittest
|
||||
|
||||
from swarms.models import OpenAIChat
|
||||
from swarms.structs import Agent, Task
|
||||
from swarms.structs.team import Team
|
||||
|
||||
|
||||
class TestTeam(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.agent = Agent(
|
||||
llm=OpenAIChat(openai_api_key=""),
|
||||
max_loops=1,
|
||||
dashboard=False,
|
||||
)
|
||||
self.task = Task(
|
||||
description="What's the weather in miami",
|
||||
agent=self.agent,
|
||||
)
|
||||
self.team = Team(
|
||||
tasks=[self.task],
|
||||
agents=[self.agent],
|
||||
architecture="sequential",
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
def test_check_config(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.team.check_config({"config": None})
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
self.team.check_config(
|
||||
{"config": json.dumps({"agents": [], "tasks": []})}
|
||||
)
|
||||
|
||||
def test_run(self):
|
||||
self.assertEqual(self.team.run(), self.task.execute())
|
||||
|
||||
def test_sequential_loop(self):
|
||||
self.assertEqual(
|
||||
self.team._Team__sequential_loop(), self.task.execute()
|
||||
)
|
||||
|
||||
def test_log(self):
|
||||
self.assertIsNone(self.team._Team__log("Test message"))
|
||||
|
||||
self.team.verbose = True
|
||||
self.assertIsNone(self.team._Team__log("Test message"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
@ -1,73 +0,0 @@
|
||||
from unittest.mock import patch
|
||||
from swarms.structs.autoscaler import AutoScaler
|
||||
from swarms.models import OpenAIChat
|
||||
from swarms.structs import Agent
|
||||
|
||||
llm = OpenAIChat()
|
||||
|
||||
agent = Agent(
|
||||
llm=llm,
|
||||
max_loops=2,
|
||||
dashboard=True,
|
||||
)
|
||||
|
||||
|
||||
def test_autoscaler_initialization():
|
||||
autoscaler = AutoScaler(
|
||||
initial_agents=5,
|
||||
scale_up_factor=2,
|
||||
idle_threshold=0.1,
|
||||
busy_threshold=0.8,
|
||||
agent=agent,
|
||||
)
|
||||
assert isinstance(autoscaler, AutoScaler)
|
||||
assert autoscaler.scale_up_factor == 2
|
||||
assert autoscaler.idle_threshold == 0.1
|
||||
assert autoscaler.busy_threshold == 0.8
|
||||
assert len(autoscaler.agents_pool) == 5
|
||||
|
||||
|
||||
def test_autoscaler_add_task():
|
||||
autoscaler = AutoScaler(agent=agent)
|
||||
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=agent
|
||||
)
|
||||
autoscaler.scale_up()
|
||||
assert len(autoscaler.agents_pool) == 10
|
||||
|
||||
|
||||
def test_autoscaler_scale_down():
|
||||
autoscaler = AutoScaler(initial_agents=5, agent=agent)
|
||||
autoscaler.scale_down()
|
||||
assert len(autoscaler.agents_pool) == 4
|
||||
|
||||
|
||||
@patch("swarms.swarms.AutoScaler.scale_up")
|
||||
@patch("swarms.swarms.AutoScaler.scale_down")
|
||||
def test_autoscaler_monitor_and_scale(mock_scale_down, mock_scale_up):
|
||||
autoscaler = AutoScaler(initial_agents=5, agent=agent)
|
||||
autoscaler.add_task("task1")
|
||||
autoscaler.monitor_and_scale()
|
||||
mock_scale_up.assert_called_once()
|
||||
mock_scale_down.assert_called_once()
|
||||
|
||||
|
||||
@patch("swarms.swarms.AutoScaler.monitor_and_scale")
|
||||
@patch("swarms.swarms.agent.run")
|
||||
def test_autoscaler_start(mock_run, mock_monitor_and_scale):
|
||||
autoscaler = AutoScaler(initial_agents=5, agent=agent)
|
||||
autoscaler.add_task("task1")
|
||||
autoscaler.start()
|
||||
mock_run.assert_called_once()
|
||||
mock_monitor_and_scale.assert_called_once()
|
||||
|
||||
|
||||
def test_autoscaler_del_agent():
|
||||
autoscaler = AutoScaler(initial_agents=5, agent=agent)
|
||||
autoscaler.del_agent()
|
||||
assert len(autoscaler.agents_pool) == 4
|
Loading…
Reference in new issue