diff --git a/swarms/models/__init__.py b/swarms/models/__init__.py index 42b18740..9339b8e2 100644 --- a/swarms/models/__init__.py +++ b/swarms/models/__init__.py @@ -46,6 +46,8 @@ from swarms.models.zeroscope import ZeroscopeTTV # noqa: E402 # from swarms.models.kosmos_two import Kosmos # noqa: E402 from swarms.models.cog_agent import CogAgent # noqa: E402 + +############## Types from swarms.models.types import ( TextModality, ImageModality, diff --git a/tests/structs/test_swarmnetwork.py b/tests/structs/test_swarmnetwork.py index 9264ee8d..683d3bb8 100644 --- a/tests/structs/test_swarmnetwork.py +++ b/tests/structs/test_swarmnetwork.py @@ -1,7 +1,10 @@ +from unittest.mock import MagicMock, Mock, patch + import pytest -from unittest.mock import Mock, patch -from swarms.structs.swarm_net import SwarmNetwork + +from swarm_net import SwarmNet from swarms.structs.agent import Agent +from swarms.structs.swarm_net import SwarmNetwork @pytest.fixture @@ -48,3 +51,42 @@ def test_swarm_network_remove_agent(swarm_network): 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") diff --git a/tests/structs/tests_graph_workflow.py b/tests/structs/tests_graph_workflow.py index bacdd681..c4bafd35 100644 --- a/tests/structs/tests_graph_workflow.py +++ b/tests/structs/tests_graph_workflow.py @@ -1,56 +1,77 @@ import pytest from swarms.structs.graph_workflow import GraphWorkflow + @pytest.fixture def graph_workflow(): return GraphWorkflow() + def test_init(graph_workflow): assert graph_workflow.graph == {} assert graph_workflow.entry_point is None + def test_add(graph_workflow): graph_workflow.add("node1", "value1") assert "node1" in graph_workflow.graph assert graph_workflow.graph["node1"]["value"] == "value1" assert graph_workflow.graph["node1"]["edges"] == {} + def test_set_entry_point(graph_workflow): graph_workflow.add("node1", "value1") graph_workflow.set_entry_point("node1") assert graph_workflow.entry_point == "node1" + def test_set_entry_point_nonexistent_node(graph_workflow): - with pytest.raises(ValueError, match="Node does not exist in graph"): + with pytest.raises( + ValueError, match="Node does not exist in graph" + ): graph_workflow.set_entry_point("nonexistent") + def test_add_edge(graph_workflow): graph_workflow.add("node1", "value1") graph_workflow.add("node2", "value2") graph_workflow.add_edge("node1", "node2") assert "node2" in graph_workflow.graph["node1"]["edges"] + def test_add_edge_nonexistent_node(graph_workflow): graph_workflow.add("node1", "value1") - with pytest.raises(ValueError, match="Node does not exist in graph"): + with pytest.raises( + ValueError, match="Node does not exist in graph" + ): graph_workflow.add_edge("node1", "nonexistent") + def test_add_conditional_edges(graph_workflow): graph_workflow.add("node1", "value1") graph_workflow.add("node2", "value2") - graph_workflow.add_conditional_edges("node1", "condition1", {"condition_value1": "node2"}) + graph_workflow.add_conditional_edges( + "node1", "condition1", {"condition_value1": "node2"} + ) assert "node2" in graph_workflow.graph["node1"]["edges"] + def test_add_conditional_edges_nonexistent_node(graph_workflow): graph_workflow.add("node1", "value1") - with pytest.raises(ValueError, match="Node does not exist in graph"): - graph_workflow.add_conditional_edges("node1", "condition1", {"condition_value1": "nonexistent"}) + with pytest.raises( + ValueError, match="Node does not exist in graph" + ): + graph_workflow.add_conditional_edges( + "node1", "condition1", {"condition_value1": "nonexistent"} + ) + def test_run(graph_workflow): graph_workflow.add("node1", "value1") graph_workflow.set_entry_point("node1") assert graph_workflow.run() == graph_workflow.graph + def test_run_no_entry_point(graph_workflow): with pytest.raises(ValueError, match="Entry point not set"): - graph_workflow.run() \ No newline at end of file + graph_workflow.run()