diff --git a/tests/models/mistral.py b/tests/models/mistral.py index 8296b106..f389cd33 100644 --- a/tests/models/mistral.py +++ b/tests/models/mistral.py @@ -2,6 +2,7 @@ import pytest from unittest.mock import patch, MagicMock from swarms.models.mistral import Mistral + def test_mistral_initialization(): mistral = Mistral(device="cpu") assert isinstance(mistral, Mistral) @@ -14,28 +15,32 @@ def test_mistral_initialization(): assert mistral.max_length == 100 assert mistral.history == [] -@patch('your_module.AutoModelForCausalLM.from_pretrained') -@patch('your_module.AutoTokenizer.from_pretrained') + +@patch("your_module.AutoModelForCausalLM.from_pretrained") +@patch("your_module.AutoTokenizer.from_pretrained") def test_mistral_load_model(mock_tokenizer, mock_model): mistral = Mistral(device="cpu") mistral.load_model() mock_model.assert_called_once() mock_tokenizer.assert_called_once() -@patch('your_module.Mistral.load_model') + +@patch("your_module.Mistral.load_model") def test_mistral_run(mock_load_model): mistral = Mistral(device="cpu") mistral.run("What's the weather in miami") mock_load_model.assert_called_once() -@patch('your_module.Mistral.run') + +@patch("your_module.Mistral.run") def test_mistral_chat(mock_run): mistral = Mistral(device="cpu") mistral.chat("What's the weather in miami") mock_run.assert_called_once() + def test_mistral__stream_response(): mistral = Mistral(device="cpu") response = "It's sunny in Miami." tokens = list(mistral._stream_response(response)) - assert tokens == ["It's", "sunny", "in", "Miami."] \ No newline at end of file + assert tokens == ["It's", "sunny", "in", "Miami."] diff --git a/tests/structs/nonlinear_workflow.py b/tests/structs/nonlinear_workflow.py index 295ec71f..863b2d9c 100644 --- a/tests/structs/nonlinear_workflow.py +++ b/tests/structs/nonlinear_workflow.py @@ -2,6 +2,7 @@ import pytest from unittest.mock import patch, MagicMock from swarms.structs.nonlinear_workflow import NonLinearWorkflow, Task + class MockTask(Task): def can_execute(self): return True @@ -9,6 +10,7 @@ class MockTask(Task): def execute(self): return "Task executed" + def test_nonlinearworkflow_initialization(): agents = MagicMock() iters_per_task = MagicMock() @@ -17,6 +19,7 @@ def test_nonlinearworkflow_initialization(): assert workflow.agents == agents assert workflow.tasks == [] + def test_nonlinearworkflow_add(): agents = MagicMock() iters_per_task = MagicMock() @@ -25,8 +28,9 @@ def test_nonlinearworkflow_add(): workflow.add(task) assert workflow.tasks == [task] -@patch('your_module.NonLinearWorkflow.is_finished') -@patch('your_module.NonLinearWorkflow.output_tasks') + +@patch("your_module.NonLinearWorkflow.is_finished") +@patch("your_module.NonLinearWorkflow.output_tasks") def test_nonlinearworkflow_run(mock_output_tasks, mock_is_finished): agents = MagicMock() iters_per_task = MagicMock() @@ -38,6 +42,7 @@ def test_nonlinearworkflow_run(mock_output_tasks, mock_is_finished): workflow.run() assert mock_output_tasks.called + def test_nonlinearworkflow_output_tasks(): agents = MagicMock() iters_per_task = MagicMock() @@ -46,6 +51,7 @@ def test_nonlinearworkflow_output_tasks(): workflow.add(task) assert workflow.output_tasks() == [task] + def test_nonlinearworkflow_to_graph(): agents = MagicMock() iters_per_task = MagicMock() @@ -54,10 +60,11 @@ def test_nonlinearworkflow_to_graph(): workflow.add(task) assert workflow.to_graph() == {"task1": set()} + def test_nonlinearworkflow_order_tasks(): agents = MagicMock() iters_per_task = MagicMock() workflow = NonLinearWorkflow(agents, iters_per_task) task = MockTask("task1") workflow.add(task) - assert workflow.order_tasks() == [task] \ No newline at end of file + assert workflow.order_tasks() == [task] diff --git a/tests/structs/workflow.py b/tests/structs/workflow.py index 5d973b18..9bab2d9d 100644 --- a/tests/structs/workflow.py +++ b/tests/structs/workflow.py @@ -2,6 +2,7 @@ import pytest from unittest.mock import patch, MagicMock from swarms.structs.workflow import Workflow + def test_workflow_initialization(): agent = MagicMock() workflow = Workflow(agent) @@ -10,6 +11,7 @@ def test_workflow_initialization(): assert workflow.tasks == [] assert workflow.parallel == False + def test_workflow_add(): agent = MagicMock() workflow = Workflow(agent) @@ -21,6 +23,7 @@ def test_workflow_add(): assert task.output == None assert task.structure == workflow + def test_workflow_first_task(): agent = MagicMock() workflow = Workflow(agent) @@ -28,6 +31,7 @@ def test_workflow_first_task(): workflow.add("What's the weather in miami") assert workflow.first_task().task == "What's the weather in miami" + def test_workflow_last_task(): agent = MagicMock() workflow = Workflow(agent) @@ -35,7 +39,8 @@ def test_workflow_last_task(): workflow.add("What's the weather in miami") assert workflow.last_task().task == "What's the weather in miami" -@patch('your_module.Workflow.__run_from_task') + +@patch("your_module.Workflow.__run_from_task") def test_workflow_run(mock_run_from_task): agent = MagicMock() workflow = Workflow(agent) @@ -43,17 +48,23 @@ def test_workflow_run(mock_run_from_task): workflow.run() mock_run_from_task.assert_called_once() + def test_workflow_context(): agent = MagicMock() workflow = Workflow(agent) task = workflow.add("What's the weather in miami") - assert workflow.context(task) == {"parent_output": None, "parent": None, "child": None} + assert workflow.context(task) == { + "parent_output": None, + "parent": None, + "child": None, + } + -@patch('your_module.Workflow.Task.execute') +@patch("your_module.Workflow.Task.execute") def test_workflow___run_from_task(mock_execute): agent = MagicMock() workflow = Workflow(agent) task = workflow.add("What's the weather in miami") mock_execute.return_value = "Sunny" workflow.__run_from_task(task) - mock_execute.assert_called_once() \ No newline at end of file + mock_execute.assert_called_once() diff --git a/tests/swarms/autoscaler.py b/tests/swarms/autoscaler.py new file mode 100644 index 00000000..5726cf5b --- /dev/null +++ b/tests/swarms/autoscaler.py @@ -0,0 +1,62 @@ +import pytest +from unittest.mock import patch, MagicMock +from swarms.swarms.autoscaler import AutoScaler, Worker + + +def test_autoscaler_initialization(): + autoscaler = AutoScaler( + initial_agents=5, + scale_up_factor=2, + idle_threshold=0.1, + busy_threshold=0.8, + agent=Worker, + ) + 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=Worker) + 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=Worker) + autoscaler.scale_up() + assert len(autoscaler.agents_pool) == 10 + + +def test_autoscaler_scale_down(): + autoscaler = AutoScaler(initial_agents=5, agent=Worker) + autoscaler.scale_down() + assert len(autoscaler.agents_pool) == 4 + + +@patch("your_module.AutoScaler.scale_up") +@patch("your_module.AutoScaler.scale_down") +def test_autoscaler_monitor_and_scale(mock_scale_down, mock_scale_up): + autoscaler = AutoScaler(initial_agents=5, agent=Worker) + autoscaler.add_task("task1") + autoscaler.monitor_and_scale() + mock_scale_up.assert_called_once() + mock_scale_down.assert_called_once() + + +@patch("your_module.AutoScaler.monitor_and_scale") +@patch("your_module.Worker.run") +def test_autoscaler_start(mock_run, mock_monitor_and_scale): + autoscaler = AutoScaler(initial_agents=5, agent=Worker) + 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=Worker) + autoscaler.del_agent() + assert len(autoscaler.agents_pool) == 4 diff --git a/tests/swarms/dialogue_simulator.py b/tests/swarms/dialogue_simulator.py new file mode 100644 index 00000000..f6bb8ad3 --- /dev/null +++ b/tests/swarms/dialogue_simulator.py @@ -0,0 +1,20 @@ +import pytest +from unittest.mock import patch, MagicMock +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.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.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 \ No newline at end of file diff --git a/tests/swarms/godmode.py b/tests/swarms/godmode.py new file mode 100644 index 00000000..44e97b25 --- /dev/null +++ b/tests/swarms/godmode.py @@ -0,0 +1,26 @@ +import pytest +from unittest.mock import patch, MagicMock +from swarms.swarms.god_mode import GodMode, LLM + +def test_godmode_initialization(): + godmode = GodMode(llms=[LLM]*5) + assert isinstance(godmode, GodMode) + assert len(godmode.llms) == 5 + +def test_godmode_run(monkeypatch): + def mock_llm_run(self, task): + return "response" + monkeypatch.setattr(LLM, "run", mock_llm_run) + godmode = GodMode(llms=[LLM]*5) + responses = godmode.run("task1") + assert len(responses) == 5 + assert responses == ["response", "response", "response", "response", "response"] + +@patch('builtins.print') +def test_godmode_print_responses(mock_print, monkeypatch): + def mock_llm_run(self, task): + return "response" + monkeypatch.setattr(LLM, "run", mock_llm_run) + godmode = GodMode(llms=[LLM]*5) + godmode.print_responses("task1") + assert mock_print.call_count == 1 \ No newline at end of file diff --git a/tests/swarms/groupchat.py b/tests/swarms/groupchat.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/swarms/multi_agent_collab.py b/tests/swarms/multi_agent_collab.py new file mode 100644 index 00000000..62b04dbf --- /dev/null +++ b/tests/swarms/multi_agent_collab.py @@ -0,0 +1,52 @@ +import pytest +from unittest.mock import patch, MagicMock +from swarms.swarms.multi_agent_collab import MultiAgentCollaboration, Worker, select_next_speaker + +def test_multiagentcollaboration_initialization(): + multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + assert isinstance(multiagentcollaboration, MultiAgentCollaboration) + assert len(multiagentcollaboration.agents) == 5 + assert multiagentcollaboration._step == 0 + +@patch('swarms.workers.Worker.reset') +def test_multiagentcollaboration_reset(mock_reset): + multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + multiagentcollaboration.reset() + assert mock_reset.call_count == 5 + +@patch('swarms.workers.Worker.run') +def test_multiagentcollaboration_inject(mock_run): + multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + multiagentcollaboration.inject("Agent 1", "Hello, world!") + assert multiagentcollaboration._step == 1 + assert mock_run.call_count == 5 + +@patch('swarms.workers.Worker.send') +@patch('swarms.workers.Worker.receive') +def test_multiagentcollaboration_step(mock_receive, mock_send): + multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + result = multiagentcollaboration.step() + assert multiagentcollaboration._step == 1 + assert mock_send.call_count == 5 + assert mock_receive.call_count == 25 + +@patch('swarms.workers.Worker.bid') +def test_multiagentcollaboration_ask_for_bid(mock_bid): + multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + result = multiagentcollaboration.ask_for_bid(Worker) + assert isinstance(result, int) + +@patch('swarms.workers.Worker.bid') +def test_multiagentcollaboration_select_next_speaker(mock_bid): + multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + result = multiagentcollaboration.select_next_speaker(1, [Worker]*5) + assert isinstance(result, int) + +@patch('swarms.workers.Worker.send') +@patch('swarms.workers.Worker.receive') +def test_multiagentcollaboration_run(mock_receive, mock_send): + multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + multiagentcollaboration.run(max_iters=5) + assert multiagentcollaboration._step == 6 + assert mock_send.call_count == 30 + assert mock_receive.call_count == 150 \ No newline at end of file diff --git a/tests/swarms/multi_agent_debate.py b/tests/swarms/multi_agent_debate.py new file mode 100644 index 00000000..f01e6a72 --- /dev/null +++ b/tests/swarms/multi_agent_debate.py @@ -0,0 +1,38 @@ +import pytest +from unittest.mock import patch, MagicMock +from swarms.swarms.multi_agent_debate import MultiAgentDebate, Worker, select_speaker + +def test_multiagentdebate_initialization(): + multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker) + assert isinstance(multiagentdebate, MultiAgentDebate) + assert len(multiagentdebate.agents) == 5 + assert multiagentdebate.selection_func == select_speaker + +@patch('swarms.workers.Worker.reset') +def test_multiagentdebate_reset_agents(mock_reset): + multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker) + multiagentdebate.reset_agents() + assert mock_reset.call_count == 5 + +def test_multiagentdebate_inject_agent(): + multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker) + multiagentdebate.inject_agent(Worker) + assert len(multiagentdebate.agents) == 6 + +@patch('swarms.workers.Worker.run') +def test_multiagentdebate_run(mock_run): + multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker) + results = multiagentdebate.run("Write a short story.") + assert len(results) == 5 + assert mock_run.call_count == 5 + +def test_multiagentdebate_update_task(): + multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker) + multiagentdebate.update_task("Write a short story.") + assert multiagentdebate.task == "Write a short story." + +def test_multiagentdebate_format_results(): + multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker) + results = [{"agent": "Agent 1", "response": "Hello, world!"}, {"agent": "Agent 2", "response": "Goodbye, world!"}] + formatted_results = multiagentdebate.format_results(results) + assert formatted_results == "Agent Agent 1 responded: Hello, world!\nAgent Agent 2 responded: Goodbye, world!" \ No newline at end of file diff --git a/tests/swarms/orchestrate.py b/tests/swarms/orchestrate.py new file mode 100644 index 00000000..af33d479 --- /dev/null +++ b/tests/swarms/orchestrate.py @@ -0,0 +1,62 @@ +import numpy as np +import pytest +from unittest.mock import patch, MagicMock +from swarms.swarms.orchestrate import Orchestrator, Worker +import chromadb + +def test_orchestrator_initialization(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + assert isinstance(orchestrator, Orchestrator) + assert orchestrator.agents.qsize() == 5 + assert orchestrator.task_queue.qsize() == 0 + +def test_orchestrator_assign_task(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + orchestrator.assign_task(1, {"content": "task1"}) + assert orchestrator.task_queue.qsize() == 1 + +def test_orchestrator_embed(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + result = orchestrator.embed("Hello, world!", "api_key", "model_name") + assert isinstance(result, np.ndarray) + +def test_orchestrator_retrieve_results(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + result = orchestrator.retrieve_results(1) + assert isinstance(result, list) + +def test_orchestrator_update_vector_db(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + data = {"vector": np.array([1, 2, 3]), "task_id": 1} + orchestrator.update_vector_db(data) + assert orchestrator.collection.count() == 1 + +def test_orchestrator_get_vector_db(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + result = orchestrator.get_vector_db() + assert isinstance(result, chromadb.Collection) + +def test_orchestrator_append_to_db(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + orchestrator.append_to_db("Hello, world!") + assert orchestrator.collection.count() == 1 + +def test_orchestrator_run(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + result = orchestrator.run("Write a short story.") + assert isinstance(result, list) + +def test_orchestrator_chat(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + orchestrator.chat(1, 2, "Hello, Agent 2!") + assert orchestrator.collection.count() == 1 + +def test_orchestrator_add_agents(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + orchestrator.add_agents(5) + assert orchestrator.agents.qsize() == 10 + +def test_orchestrator_remove_agents(): + orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[]) + orchestrator.remove_agents(3) + assert orchestrator.agents.qsize() == 2 \ No newline at end of file diff --git a/tests/swarms/scalable_groupchat.py b/tests/swarms/scalable_groupchat.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/swarms/simple_swarm.py b/tests/swarms/simple_swarm.py new file mode 100644 index 00000000..4753a99b --- /dev/null +++ b/tests/swarms/simple_swarm.py @@ -0,0 +1,37 @@ +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.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 \ No newline at end of file