[RMEOVING TESTS FOR ORCHESTRATE, DIALOGUE SIMULATOR etc]

pull/286/head
Kye 1 year ago
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…
Cancel
Save