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.
93 lines
2.4 KiB
93 lines
2.4 KiB
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import pytest
|
|
|
|
from swarm_net import SwarmNet
|
|
from swarms.structs.agent import Agent
|
|
from swarms.structs.swarm_net import SwarmNetwork
|
|
|
|
|
|
@pytest.fixture
|
|
def swarm_network():
|
|
agents = [Agent(id=f"Agent_{i}") for i in range(5)]
|
|
return SwarmNetwork(agents=agents)
|
|
|
|
|
|
def test_swarm_network_init(swarm_network):
|
|
assert isinstance(swarm_network.agents, list)
|
|
assert len(swarm_network.agents) == 5
|
|
|
|
|
|
@patch("swarms.structs.swarm_net.SwarmNetwork.logger")
|
|
def test_run(mock_logger, swarm_network):
|
|
swarm_network.run()
|
|
assert (
|
|
mock_logger.info.call_count == 10
|
|
) # 2 log messages per agent
|
|
|
|
|
|
def test_run_with_mocked_agents(mocker, swarm_network):
|
|
mock_agents = [Mock(spec=Agent) for _ in range(5)]
|
|
mocker.patch.object(swarm_network, "agents", mock_agents)
|
|
swarm_network.run()
|
|
for mock_agent in mock_agents:
|
|
assert mock_agent.run.called
|
|
|
|
|
|
def test_swarm_network_with_no_agents():
|
|
swarm_network = SwarmNetwork(agents=[])
|
|
assert swarm_network.agents == []
|
|
|
|
|
|
def test_swarm_network_add_agent(swarm_network):
|
|
new_agent = Agent(id="Agent_5")
|
|
swarm_network.add_agent(new_agent)
|
|
assert len(swarm_network.agents) == 6
|
|
assert swarm_network.agents[-1] == new_agent
|
|
|
|
|
|
def test_swarm_network_remove_agent(swarm_network):
|
|
agent_to_remove = swarm_network.agents[0]
|
|
swarm_network.remove_agent(agent_to_remove)
|
|
assert len(swarm_network.agents) == 4
|
|
assert agent_to_remove not in swarm_network.agents
|
|
|
|
|
|
@pytest.fixture
|
|
def swarmnet():
|
|
swarmnet = SwarmNet()
|
|
agent_mock = MagicMock()
|
|
agent_mock.id = "1"
|
|
swarmnet.agents = [agent_mock]
|
|
return swarmnet
|
|
|
|
|
|
def test_run_agent(swarmnet):
|
|
swarmnet.run_agent("1", "task")
|
|
swarmnet.agents[0].run.assert_called_once_with("task")
|
|
|
|
|
|
def test_run_agent_no_agent(swarmnet):
|
|
with pytest.raises(ValueError, match="No agent found with ID"):
|
|
swarmnet.run_agent("2", "task")
|
|
|
|
|
|
def test_run_many_agents(swarmnet):
|
|
swarmnet.run_many_agents("task")
|
|
swarmnet.agents[0].run.assert_called_once_with("task")
|
|
|
|
|
|
def test_list_agents(swarmnet):
|
|
swarmnet.list_agents()
|
|
assert swarmnet.agents[0].id == "1"
|
|
|
|
|
|
def test_get_agent(swarmnet):
|
|
agent = swarmnet.get_agent("1")
|
|
assert agent.id == "1"
|
|
|
|
|
|
def test_get_agent_no_agent(swarmnet):
|
|
with pytest.raises(ValueError, match="No agent found with ID"):
|
|
swarmnet.get_agent("2")
|