You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
swarms/tests/swarms/simple_swarm.py

52 lines
1.5 KiB

from unittest.mock import patch
from swarms.swarms.simple_swarm import SimpleSwarm
def test_simpleswarm_initialization():
simpleswarm = SimpleSwarm(
num_workers=5, openai_api_key="api_key", ai_name="ai_name"
)
assert isinstance(simpleswarm, SimpleSwarm)
assert len(simpleswarm.workers) == 5
assert simpleswarm.task_queue.qsize() == 0
assert simpleswarm.priority_queue.qsize() == 0
def test_simpleswarm_distribute():
simpleswarm = SimpleSwarm(
num_workers=5, openai_api_key="api_key", ai_name="ai_name"
)
simpleswarm.distribute("task1")
assert simpleswarm.task_queue.qsize() == 1
simpleswarm.distribute("task2", priority=1)
assert simpleswarm.priority_queue.qsize() == 1
@patch("swarms.workers.worker.Worker.run")
def test_simpleswarm_process_task(mock_run):
simpleswarm = SimpleSwarm(
num_workers=5, openai_api_key="api_key", ai_name="ai_name"
)
simpleswarm._process_task("task1")
assert mock_run.call_count == 5
def test_simpleswarm_run():
simpleswarm = SimpleSwarm(
num_workers=5, openai_api_key="api_key", ai_name="ai_name"
)
simpleswarm.distribute("task1")
simpleswarm.distribute("task2", priority=1)
results = simpleswarm.run()
assert len(results) == 2
@patch("swarms.workers.Worker.run")
def test_simpleswarm_run_old(mock_run):
simpleswarm = SimpleSwarm(
num_workers=5, openai_api_key="api_key", ai_name="ai_name"
)
results = simpleswarm.run_old("task1")
assert len(results) == 5
assert mock_run.call_count == 5