53 lines
1.6 KiB
53 lines
1.6 KiB
import pytest
|
|
from unittest.mock import patch, MagicMock
|
|
from swarms.swarms.simple_swarm import SimpleSwarm, Worker
|
|
|
|
|
|
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"
|
|
)
|
|
result = 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
|