[TESTS][SwarmNetwork]

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,

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")

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"):
+    with pytest.raises(
-        graph_workflow.add_conditional_edges("node1", "condition1", {"condition_value1": "nonexistent"})
+        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()
+        graph_workflow.run()