parent
74d2a2fde9
commit
7529c5280e
@ -1,24 +0,0 @@
|
||||
from unittest.mock import patch
|
||||
from swarms.swarms.dialogue_simulator import DialogueSimulator, Worker
|
||||
|
||||
|
||||
def test_dialoguesimulator_initialization():
|
||||
dialoguesimulator = DialogueSimulator(agents=[Worker] * 5)
|
||||
assert isinstance(dialoguesimulator, DialogueSimulator)
|
||||
assert len(dialoguesimulator.agents) == 5
|
||||
|
||||
|
||||
@patch("swarms.workers.worker.Worker.run")
|
||||
def test_dialoguesimulator_run(mock_run):
|
||||
dialoguesimulator = DialogueSimulator(agents=[Worker] * 5)
|
||||
dialoguesimulator.run(
|
||||
max_iters=5, name="Agent 1", message="Hello, world!"
|
||||
)
|
||||
assert mock_run.call_count == 30
|
||||
|
||||
|
||||
@patch("swarms.workers.worker.Worker.run")
|
||||
def test_dialoguesimulator_run_without_name_and_message(mock_run):
|
||||
dialoguesimulator = DialogueSimulator(agents=[Worker] * 5)
|
||||
dialoguesimulator.run(max_iters=5)
|
||||
assert mock_run.call_count == 25
|
@ -1,72 +0,0 @@
|
||||
import pytest
|
||||
from unittest.mock import Mock
|
||||
from swarms.swarms.orchestrate import Orchestrator
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent():
|
||||
return Mock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_task():
|
||||
return {"task_id": 1, "task_data": "data"}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_vector_db():
|
||||
return Mock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def orchestrator(mock_agent, mock_vector_db):
|
||||
agent_list = [mock_agent for _ in range(5)]
|
||||
task_queue = []
|
||||
return Orchestrator(
|
||||
mock_agent, agent_list, task_queue, mock_vector_db
|
||||
)
|
||||
|
||||
|
||||
def test_assign_task(
|
||||
orchestrator, mock_agent, mock_task, mock_vector_db
|
||||
):
|
||||
orchestrator.task_queue.append(mock_task)
|
||||
orchestrator.assign_task(0, mock_task)
|
||||
|
||||
mock_agent.process_task.assert_called_once()
|
||||
mock_vector_db.add_documents.assert_called_once()
|
||||
|
||||
|
||||
def test_retrieve_results(orchestrator, mock_vector_db):
|
||||
mock_vector_db.query.return_value = "expected_result"
|
||||
assert orchestrator.retrieve_results(0) == "expected_result"
|
||||
|
||||
|
||||
def test_update_vector_db(orchestrator, mock_vector_db):
|
||||
data = {"vector": [0.1, 0.2, 0.3], "task_id": 1}
|
||||
orchestrator.update_vector_db(data)
|
||||
mock_vector_db.add_documents.assert_called_once_with(
|
||||
[data["vector"]], [str(data["task_id"])]
|
||||
)
|
||||
|
||||
|
||||
def test_get_vector_db(orchestrator, mock_vector_db):
|
||||
assert orchestrator.get_vector_db() == mock_vector_db
|
||||
|
||||
|
||||
def test_append_to_db(orchestrator, mock_vector_db):
|
||||
collection = "test_collection"
|
||||
result = "test_result"
|
||||
orchestrator.append_to_db(collection, result)
|
||||
mock_vector_db.append_document.assert_called_once_with(
|
||||
collection, result, id=str(id(result))
|
||||
)
|
||||
|
||||
|
||||
def test_run(orchestrator, mock_agent, mock_vector_db):
|
||||
objective = "test_objective"
|
||||
collection = "test_collection"
|
||||
orchestrator.run(objective, collection)
|
||||
|
||||
mock_agent.process_task.assert_called()
|
||||
mock_vector_db.append_document.assert_called()
|
@ -1,51 +0,0 @@
|
||||
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
|
Loading…
Reference in new issue