[test][swarm router]

tests/requirements.txt

@@ -2,3 +2,4 @@ swarms
 pytest
 matplotlib
 loguru
+unittest

tests/structs/test_swarm_router.py

@@ -1,34 +1,18 @@
-"""
-Comprehensive unit test suite for SwarmRouter class.
-
-This test suite covers all functions, swarm types, and ensures outputs are not None
-to verify successful execution.
-"""
-
 import pytest
-import asyncio
-from unittest.mock import Mock, patch, MagicMock
-from typing import List, Dict, Any
-import concurrent.futures
-import json
+from unittest.mock import Mock, patch
 
 from swarms.structs.swarm_router import (
     SwarmRouter,
     SwarmRouterConfig,
     SwarmRouterRunError,
     SwarmRouterConfigError,
-    SwarmType,
     Document,
 )
 from swarms.structs.agent import Agent
-from swarms.utils.output_types import OutputType
-
 
-class TestSwarmRouterInitialization:
-    """Test SwarmRouter initialization with various configurations."""
 
 @pytest.fixture
-    def sample_agents(self):
+def sample_agents():
     """Create sample agents for testing."""
     return [
         Agent(
@@ -45,12 +29,15 @@ class TestSwarmRouterInitialization:
         ),
     ]
 
-    def test_default_initialization(self):
+
+def test_default_initialization():
     """Test SwarmRouter with default parameters."""
     router = SwarmRouter()
 
     assert router.name == "swarm-router"
-        assert router.description == "Routes your task to the desired swarm"
+    assert (
+        router.description == "Routes your task to the desired swarm"
+    )
     assert router.max_loops == 1
     assert router.agents == []
     assert router.swarm_type == "SequentialWorkflow"
@@ -64,7 +51,8 @@ class TestSwarmRouterInitialization:
     assert router.verbose is False
     assert router.telemetry_enabled is False
 
-    def test_custom_initialization(self, sample_agents):
+
+def test_custom_initialization(sample_agents):
     """Test SwarmRouter with custom parameters."""
     router = SwarmRouter(
         name="test-router",
@@ -96,7 +84,8 @@ class TestSwarmRouterInitialization:
     assert router.verbose is True
     assert router.telemetry_enabled is True
 
-    def test_initialization_with_heavy_swarm_config(self, sample_agents):
+
+def test_initialization_with_heavy_swarm_config(sample_agents):
     """Test SwarmRouter with HeavySwarm specific configuration."""
     router = SwarmRouter(
         agents=sample_agents,
@@ -113,7 +102,8 @@ class TestSwarmRouterInitialization:
     assert router.heavy_swarm_worker_model_name == "gpt-3.5-turbo"
     assert router.heavy_swarm_swarm_show_output is False
 
-    def test_initialization_with_council_judge_config(self):
+
+def test_initialization_with_council_judge_config():
     """Test SwarmRouter with CouncilAsAJudge specific configuration."""
     router = SwarmRouter(
         swarm_type="CouncilAsAJudge",
@@ -123,7 +113,8 @@ class TestSwarmRouterInitialization:
     assert router.swarm_type == "CouncilAsAJudge"
     assert router.council_judge_model_name == "gpt-4o"
 
-    def test_initialization_with_agent_rearrange_flow(self, sample_agents):
+
+def test_initialization_with_agent_rearrange_flow(sample_agents):
     """Test SwarmRouter with AgentRearrange and flow configuration."""
     flow = "agent1 -> agent2 -> agent1"
     router = SwarmRouter(
@@ -135,7 +126,8 @@ class TestSwarmRouterInitialization:
     assert router.swarm_type == "AgentRearrange"
     assert router.rearrange_flow == flow
 
-    def test_initialization_with_shared_memory(self, sample_agents):
+
+def test_initialization_with_shared_memory(sample_agents):
     """Test SwarmRouter with shared memory system."""
     mock_memory = Mock()
     router = SwarmRouter(
@@ -145,7 +137,8 @@ class TestSwarmRouterInitialization:
 
     assert router.shared_memory_system == mock_memory
 
-    def test_initialization_with_worker_tools(self, sample_agents):
+
+def test_initialization_with_worker_tools(sample_agents):
     """Test SwarmRouter with worker tools."""
     mock_tool = Mock()
     router = SwarmRouter(
@@ -156,20 +149,19 @@ class TestSwarmRouterInitialization:
     assert router.worker_tools == [mock_tool]
 
 
-class TestSwarmRouterConfigurationValidation:
-    """Test configuration validation and error handling."""
-
-    def test_invalid_swarm_type(self):
+def test_invalid_swarm_type():
     """Test error when invalid swarm type is provided."""
     with pytest.raises(ValueError):
         SwarmRouter(swarm_type="InvalidSwarmType")
 
-    def test_no_agents_for_swarm_requiring_agents(self):
+
+def test_no_agents_for_swarm_requiring_agents():
     """Test error when no agents provided for swarm requiring agents."""
     with pytest.raises(SwarmRouterConfigError):
         SwarmRouter(swarm_type="SequentialWorkflow", agents=None)
 
-    def test_no_rearrange_flow_for_agent_rearrange(self):
+
+def test_no_rearrange_flow_for_agent_rearrange():
     """Test error when no rearrange_flow provided for AgentRearrange."""
     agents = [Agent(agent_name="test", agent_description="test")]
     with pytest.raises(SwarmRouterConfigError):
@@ -179,60 +171,52 @@ class TestSwarmRouterConfigurationValidation:
             rearrange_flow=None,
         )
 
-    def test_zero_max_loops(self):
+
+def test_zero_max_loops():
     """Test error when max_loops is 0."""
     with pytest.raises(SwarmRouterConfigError):
         SwarmRouter(max_loops=0)
 
-    def test_heavy_swarm_without_agents(self):
+
+def test_heavy_swarm_without_agents():
     """Test HeavySwarm can be created without agents."""
     router = SwarmRouter(swarm_type="HeavySwarm", agents=None)
     assert router.swarm_type == "HeavySwarm"
 
-    def test_council_judge_without_agents(self):
+
+def test_council_judge_without_agents():
     """Test CouncilAsAJudge can be created without agents."""
     router = SwarmRouter(swarm_type="CouncilAsAJudge", agents=None)
     assert router.swarm_type == "CouncilAsAJudge"
 
 
-class TestSwarmFactoryMethods:
-    """Test all factory methods for creating different swarm types."""
-
-    @pytest.fixture
-    def sample_agents(self):
-        """Create sample agents for testing."""
-        return [
-            Agent(
-                agent_name="Agent1",
-                agent_description="First agent",
-                system_prompt="You are agent 1.",
-                max_loops=1,
-            ),
-            Agent(
-                agent_name="Agent2",
-                agent_description="Second agent",
-                system_prompt="You are agent 2.",
-                max_loops=1,
-            ),
-        ]
-
-    def test_swarm_factory_initialization(self, sample_agents):
+def test_swarm_factory_initialization(sample_agents):
     """Test that swarm factory is properly initialized."""
     router = SwarmRouter(agents=sample_agents)
     factory = router._initialize_swarm_factory()
 
     expected_types = [
-            "HeavySwarm", "AgentRearrange", "MALT", "CouncilAsAJudge",
-            "InteractiveGroupChat", "HiearchicalSwarm", "MixtureOfAgents",
-            "MajorityVoting", "GroupChat", "MultiAgentRouter",
-            "SequentialWorkflow", "ConcurrentWorkflow", "BatchedGridWorkflow"
+        "HeavySwarm",
+        "AgentRearrange",
+        "MALT",
+        "CouncilAsAJudge",
+        "InteractiveGroupChat",
+        "HiearchicalSwarm",
+        "MixtureOfAgents",
+        "MajorityVoting",
+        "GroupChat",
+        "MultiAgentRouter",
+        "SequentialWorkflow",
+        "ConcurrentWorkflow",
+        "BatchedGridWorkflow",
     ]
 
     for swarm_type in expected_types:
         assert swarm_type in factory
         assert callable(factory[swarm_type])
 
-    def test_create_heavy_swarm(self, sample_agents):
+
+def test_create_heavy_swarm(sample_agents):
     """Test HeavySwarm creation."""
     router = SwarmRouter(
         agents=sample_agents,
@@ -244,9 +228,10 @@ class TestSwarmFactoryMethods:
 
     swarm = router._create_heavy_swarm()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
+
 
-    def test_create_agent_rearrange(self, sample_agents):
+def test_create_agent_rearrange(sample_agents):
     """Test AgentRearrange creation."""
     router = SwarmRouter(
         agents=sample_agents,
@@ -256,119 +241,127 @@ class TestSwarmFactoryMethods:
 
     swarm = router._create_agent_rearrange()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
 
-    def test_create_sequential_workflow(self, sample_agents):
+
+def test_create_sequential_workflow(sample_agents):
     """Test SequentialWorkflow creation."""
-        router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
+    router = SwarmRouter(
+        agents=sample_agents, swarm_type="SequentialWorkflow"
+    )
 
     swarm = router._create_sequential_workflow()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
+
 
-    def test_create_concurrent_workflow(self, sample_agents):
+def test_create_concurrent_workflow(sample_agents):
     """Test ConcurrentWorkflow creation."""
-        router = SwarmRouter(agents=sample_agents, swarm_type="ConcurrentWorkflow")
+    router = SwarmRouter(
+        agents=sample_agents, swarm_type="ConcurrentWorkflow"
+    )
 
     swarm = router._create_concurrent_workflow()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
+
 
-    def test_create_group_chat(self, sample_agents):
+def test_create_group_chat(sample_agents):
     """Test GroupChat creation."""
     router = SwarmRouter(agents=sample_agents, swarm_type="GroupChat")
 
     swarm = router._create_group_chat()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
+
 
-    def test_create_multi_agent_router(self, sample_agents):
+def test_create_multi_agent_router(sample_agents):
     """Test MultiAgentRouter creation."""
-        router = SwarmRouter(agents=sample_agents, swarm_type="MultiAgentRouter")
+    router = SwarmRouter(
+        agents=sample_agents, swarm_type="MultiAgentRouter"
+    )
 
     swarm = router._create_multi_agent_router()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
+
 
-    def test_create_mixture_of_agents(self, sample_agents):
+def test_create_mixture_of_agents(sample_agents):
     """Test MixtureOfAgents creation."""
-        router = SwarmRouter(agents=sample_agents, swarm_type="MixtureOfAgents")
+    router = SwarmRouter(
+        agents=sample_agents, swarm_type="MixtureOfAgents"
+    )
 
     swarm = router._create_mixture_of_agents()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
 
-    def test_create_majority_voting(self, sample_agents):
+
+def test_create_majority_voting(sample_agents):
     """Test MajorityVoting creation."""
-        router = SwarmRouter(agents=sample_agents, swarm_type="MajorityVoting")
+    router = SwarmRouter(
+        agents=sample_agents, swarm_type="MajorityVoting"
+    )
 
     swarm = router._create_majority_voting()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
 
-    def test_create_malt(self, sample_agents):
+
+def test_create_malt(sample_agents):
     """Test MALT creation."""
     router = SwarmRouter(agents=sample_agents, swarm_type="MALT")
 
     swarm = router._create_malt()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
+
 
-    def test_create_council_as_judge(self):
+def test_create_council_as_judge():
     """Test CouncilAsAJudge creation."""
     router = SwarmRouter(swarm_type="CouncilAsAJudge")
 
     swarm = router._create_council_as_judge()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
+
 
-    def test_create_interactive_group_chat(self, sample_agents):
+def test_create_interactive_group_chat(sample_agents):
     """Test InteractiveGroupChat creation."""
-        router = SwarmRouter(agents=sample_agents, swarm_type="InteractiveGroupChat")
+    router = SwarmRouter(
+        agents=sample_agents, swarm_type="InteractiveGroupChat"
+    )
 
     swarm = router._create_interactive_group_chat()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
+
 
-    def test_create_hierarchical_swarm(self, sample_agents):
+def test_create_hierarchical_swarm(sample_agents):
     """Test HierarchicalSwarm creation."""
-        router = SwarmRouter(agents=sample_agents, swarm_type="HiearchicalSwarm")
+    router = SwarmRouter(
+        agents=sample_agents, swarm_type="HiearchicalSwarm"
+    )
 
     swarm = router._create_hierarchical_swarm()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
 
-    def test_create_batched_grid_workflow(self, sample_agents):
+
+def test_create_batched_grid_workflow(sample_agents):
     """Test BatchedGridWorkflow creation."""
-        router = SwarmRouter(agents=sample_agents, swarm_type="BatchedGridWorkflow")
+    router = SwarmRouter(
+        agents=sample_agents, swarm_type="BatchedGridWorkflow"
+    )
 
     swarm = router._create_batched_grid_workflow()
     assert swarm is not None
-        assert hasattr(swarm, 'run')
+    assert hasattr(swarm, "run")
 
 
-class TestSwarmRouterSwarmTypes:
-    """Test all available swarm types with actual execution."""
-
-    @pytest.fixture
-    def sample_agents(self):
-        """Create sample agents for testing."""
-        return [
-            Agent(
-                agent_name="ResearchAgent",
-                agent_description="Specializes in researching topics",
-                system_prompt="You are a research specialist. Provide detailed, well-researched information.",
-                max_loops=1,
-            ),
-            Agent(
-                agent_name="CodeAgent",
-                agent_description="Expert in coding",
-                system_prompt="You are a coding expert. Write clean, efficient code.",
-                max_loops=1,
-            ),
-        ]
-
-    @pytest.mark.parametrize("swarm_type", [
+@pytest.mark.parametrize(
+    "swarm_type",
+    [
         "SequentialWorkflow",
         "ConcurrentWorkflow",
         "GroupChat",
@@ -380,13 +373,16 @@ class TestSwarmRouterSwarmTypes:
         "InteractiveGroupChat",
         "HiearchicalSwarm",
         "BatchedGridWorkflow",
-    ])
-    def test_swarm_types_execution(self, sample_agents, swarm_type):
+    ],
+)
+def test_swarm_types_execution(sample_agents, swarm_type):
     """Test execution of all swarm types with mock LLM."""
-        with patch('swarms.structs.agent.LiteLLM') as mock_llm:
+    with patch("swarms.structs.agent.LiteLLM") as mock_llm:
         # Mock the LLM to return a simple response
         mock_llm_instance = Mock()
-            mock_llm_instance.agenerate.return_value = "Test response from agent"
+        mock_llm_instance.agenerate.return_value = (
+            "Test response from agent"
+        )
         mock_llm.return_value = mock_llm_instance
 
         router = SwarmRouter(
@@ -400,8 +396,10 @@ class TestSwarmRouterSwarmTypes:
 
         try:
             result = router.run(task)
-                assert result is not None, f"Swarm type {swarm_type} returned None result"
-            except Exception as e:
+            assert (
+                result is not None
+            ), f"Swarm type {swarm_type} returned None result"
+        except Exception:
             # Some swarm types might have specific requirements
             if swarm_type in ["AgentRearrange"]:
                 # AgentRearrange requires rearrange_flow
@@ -412,11 +410,16 @@ class TestSwarmRouterSwarmTypes:
                     max_loops=1,
                 )
                 result = router.run(task)
-                    assert result is not None, f"Swarm type {swarm_type} returned None result"
+                assert (
+                    result is not None
+                ), f"Swarm type {swarm_type} returned None result"
 
-    def test_heavy_swarm_execution(self):
+
+def test_heavy_swarm_execution():
     """Test HeavySwarm execution."""
-        with patch('swarms.structs.heavy_swarm.HeavySwarm') as mock_heavy_swarm:
+    with patch(
+        "swarms.structs.heavy_swarm.HeavySwarm"
+    ) as mock_heavy_swarm:
         mock_instance = Mock()
         mock_instance.run.return_value = "HeavySwarm response"
         mock_heavy_swarm.return_value = mock_instance
@@ -427,9 +430,12 @@ class TestSwarmRouterSwarmTypes:
         assert result is not None
         assert result == "HeavySwarm response"
 
-    def test_agent_rearrange_execution(self, sample_agents):
+
+def test_agent_rearrange_execution(sample_agents):
     """Test AgentRearrange execution with flow."""
-        with patch('swarms.structs.agent_rearrange.AgentRearrange') as mock_agent_rearrange:
+    with patch(
+        "swarms.structs.agent_rearrange.AgentRearrange"
+    ) as mock_agent_rearrange:
         mock_instance = Mock()
         mock_instance.run.return_value = "AgentRearrange response"
         mock_agent_rearrange.return_value = mock_instance
@@ -445,69 +451,75 @@ class TestSwarmRouterSwarmTypes:
         assert result == "AgentRearrange response"
 
 
-class TestSwarmRouterRunMethods:
-    """Test run, batch_run, concurrent_run methods."""
-
-    @pytest.fixture
-    def sample_agents(self):
-        """Create sample agents for testing."""
-        return [
-            Agent(
-                agent_name="TestAgent",
-                agent_description="Test agent",
-                system_prompt="You are a test agent.",
-                max_loops=1,
-            ),
-        ]
-
-    def test_run_method(self, sample_agents):
+def test_run_method(sample_agents):
     """Test basic run method."""
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = "Test response"
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
+        router = SwarmRouter(
+            agents=sample_agents, swarm_type="SequentialWorkflow"
+        )
 
         result = router.run("Test task")
         assert result is not None
         assert result == "Test response"
 
-    def test_run_with_image(self, sample_agents):
+
+def test_run_with_image(sample_agents):
     """Test run method with image input."""
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = "Image analysis response"
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
+        router = SwarmRouter(
+            agents=sample_agents, swarm_type="SequentialWorkflow"
+        )
 
-            result = router.run("Analyze this image", img="test_image.jpg")
+        result = router.run(
+            "Analyze this image", img="test_image.jpg"
+        )
         assert result is not None
         assert result == "Image analysis response"
 
-    def test_run_with_tasks_list(self, sample_agents):
+
+def test_run_with_tasks_list(sample_agents):
     """Test run method with tasks list."""
-        with patch('swarms.structs.batched_grid_workflow.BatchedGridWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.batched_grid_workflow.BatchedGridWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = "Batch response"
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=sample_agents, swarm_type="BatchedGridWorkflow")
+        router = SwarmRouter(
+            agents=sample_agents, swarm_type="BatchedGridWorkflow"
+        )
 
         tasks = ["Task 1", "Task 2", "Task 3"]
         result = router.run(tasks=tasks)
         assert result is not None
         assert result == "Batch response"
 
-    def test_batch_run_method(self, sample_agents):
+
+def test_batch_run_method(sample_agents):
     """Test batch_run method."""
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = "Batch task response"
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
+        router = SwarmRouter(
+            agents=sample_agents, swarm_type="SequentialWorkflow"
+        )
 
         tasks = ["Task 1", "Task 2", "Task 3"]
         results = router.batch_run(tasks)
@@ -516,75 +528,80 @@ class TestSwarmRouterRunMethods:
         assert len(results) == 3
         assert all(result is not None for result in results)
 
-    def test_concurrent_run_method(self, sample_agents):
+
+def test_concurrent_run_method(sample_agents):
     """Test concurrent_run method."""
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = "Concurrent response"
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
+        router = SwarmRouter(
+            agents=sample_agents, swarm_type="SequentialWorkflow"
+        )
 
         result = router.concurrent_run("Test task")
         assert result is not None
         assert result == "Concurrent response"
 
-    def test_call_method(self, sample_agents):
+
+def test_call_method(sample_agents):
     """Test __call__ method."""
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = "Call response"
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
+        router = SwarmRouter(
+            agents=sample_agents, swarm_type="SequentialWorkflow"
+        )
 
         result = router("Test task")
         assert result is not None
         assert result == "Call response"
 
-    def test_call_with_image(self, sample_agents):
+
+def test_call_with_image(sample_agents):
     """Test __call__ method with image."""
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = "Image call response"
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
+        router = SwarmRouter(
+            agents=sample_agents, swarm_type="SequentialWorkflow"
+        )
 
         result = router("Test task", img="test.jpg")
         assert result is not None
         assert result == "Image call response"
 
-    def test_call_with_images_list(self, sample_agents):
+
+def test_call_with_images_list(sample_agents):
     """Test __call__ method with images list."""
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = "Images call response"
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=sample_agents, swarm_type="SequentialWorkflow")
+        router = SwarmRouter(
+            agents=sample_agents, swarm_type="SequentialWorkflow"
+        )
 
         result = router("Test task", imgs=["test1.jpg", "test2.jpg"])
         assert result is not None
         assert result == "Images call response"
 
 
-class TestSwarmRouterUtilityMethods:
-    """Test utility methods like to_dict, activate_ape, handle_rules, etc."""
-
-    @pytest.fixture
-    def sample_agents(self):
-        """Create sample agents for testing."""
-        return [
-            Agent(
-                agent_name="TestAgent",
-                agent_description="Test agent",
-                system_prompt="You are a test agent.",
-                max_loops=1,
-            ),
-        ]
-
-    def test_to_dict_method(self, sample_agents):
+def test_to_dict_method(sample_agents):
     """Test to_dict method serialization."""
     router = SwarmRouter(
         agents=sample_agents,
@@ -601,7 +618,8 @@ class TestSwarmRouterUtilityMethods:
     assert result_dict["swarm_type"] == "SequentialWorkflow"
     assert "agents" in result_dict
 
-    def test_activate_ape(self, sample_agents):
+
+def test_activate_ape(sample_agents):
     """Test activate_ape method."""
     router = SwarmRouter(
         agents=sample_agents,
@@ -619,7 +637,8 @@ class TestSwarmRouterUtilityMethods:
         if hasattr(agent, "auto_generate_prompt"):
             assert agent.auto_generate_prompt is True
 
-    def test_handle_rules(self, sample_agents):
+
+def test_handle_rules(sample_agents):
     """Test handle_rules method."""
     rules = "Always be helpful and accurate."
     router = SwarmRouter(
@@ -627,15 +646,21 @@ class TestSwarmRouterUtilityMethods:
         rules=rules,
     )
 
-        original_prompts = [agent.system_prompt for agent in router.agents]
+    original_prompts = [
+        agent.system_prompt for agent in router.agents
+    ]
     router.handle_rules()
 
     # Check that rules were added to system prompts
     for i, agent in enumerate(router.agents):
         assert rules in agent.system_prompt
-            assert agent.system_prompt == original_prompts[i] + f"### Swarm Rules ### {rules}"
+        assert (
+            agent.system_prompt
+            == original_prompts[i] + f"### Swarm Rules ### {rules}"
+        )
+
 
-    def test_activate_shared_memory(self, sample_agents):
+def test_activate_shared_memory(sample_agents):
     """Test activate_shared_memory method."""
     mock_memory = Mock()
     router = SwarmRouter(
@@ -649,21 +674,25 @@ class TestSwarmRouterUtilityMethods:
     for agent in router.agents:
         assert agent.long_term_memory == mock_memory
 
-    def test_update_system_prompt_for_agent_in_swarm(self, sample_agents):
+
+def test_update_system_prompt_for_agent_in_swarm(sample_agents):
     """Test update_system_prompt_for_agent_in_swarm method."""
     router = SwarmRouter(
         agents=sample_agents,
         multi_agent_collab_prompt=True,
     )
 
-        original_prompts = [agent.system_prompt for agent in router.agents]
+    original_prompts = [
+        agent.system_prompt for agent in router.agents
+    ]
     router.update_system_prompt_for_agent_in_swarm()
 
     # Check that collaboration prompt was added
     for i, agent in enumerate(router.agents):
         assert len(agent.system_prompt) > len(original_prompts[i])
 
-    def test_agent_config(self, sample_agents):
+
+def test_agent_config(sample_agents):
     """Test agent_config method."""
     router = SwarmRouter(agents=sample_agents)
 
@@ -675,70 +704,86 @@ class TestSwarmRouterUtilityMethods:
     for agent in sample_agents:
         assert agent.agent_name in config
 
-    def test_fetch_message_history_as_string(self, sample_agents):
+
+def test_fetch_message_history_as_string(sample_agents):
     """Test fetch_message_history_as_string method."""
     router = SwarmRouter(agents=sample_agents)
 
     # Mock the swarm and conversation
     mock_conversation = Mock()
-        mock_conversation.return_all_except_first_string.return_value = "Test history"
+    mock_conversation.return_all_except_first_string.return_value = (
+        "Test history"
+    )
     router.swarm = Mock()
     router.swarm.conversation = mock_conversation
 
     result = router.fetch_message_history_as_string()
     assert result == "Test history"
 
-    def test_fetch_message_history_as_string_error(self, sample_agents):
+
+def test_fetch_message_history_as_string_error(sample_agents):
     """Test fetch_message_history_as_string method with error."""
     router = SwarmRouter(agents=sample_agents)
 
     # Mock the swarm to raise an exception
     router.swarm = Mock()
-        router.swarm.conversation.return_all_except_first_string.side_effect = Exception("Test error")
+    router.swarm.conversation.return_all_except_first_string.side_effect = Exception(
+        "Test error"
+    )
 
     result = router.fetch_message_history_as_string()
     assert result is None
 
 
-class TestSwarmRouterErrorHandling:
-    """Test error handling and exception scenarios."""
-
-    def test_swarm_creation_error(self):
+def test_swarm_creation_error():
     """Test error handling when swarm creation fails."""
     router = SwarmRouter(swarm_type="SequentialWorkflow", agents=None)
 
     with pytest.raises(SwarmRouterConfigError):
         router.run("Test task")
 
-    def test_run_error_handling(self):
+
+def test_run_error_handling():
     """Test error handling during task execution."""
     agents = [Agent(agent_name="test", agent_description="test")]
 
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
-            mock_instance.run.side_effect = Exception("Test execution error")
+        mock_instance.run.side_effect = Exception(
+            "Test execution error"
+        )
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=agents, swarm_type="SequentialWorkflow")
+        router = SwarmRouter(
+            agents=agents, swarm_type="SequentialWorkflow"
+        )
 
         with pytest.raises(SwarmRouterRunError):
             router.run("Test task")
 
-    def test_batch_run_error_handling(self):
+
+def test_batch_run_error_handling():
     """Test error handling during batch execution."""
     agents = [Agent(agent_name="test", agent_description="test")]
 
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.side_effect = Exception("Test batch error")
         mock_workflow.return_value = mock_instance
 
-            router = SwarmRouter(agents=agents, swarm_type="SequentialWorkflow")
+        router = SwarmRouter(
+            agents=agents, swarm_type="SequentialWorkflow"
+        )
 
         with pytest.raises(RuntimeError):
             router.batch_run(["Task 1", "Task 2"])
 
-    def test_invalid_swarm_type_error(self):
+
+def test_invalid_swarm_type_error():
     """Test error when creating swarm with invalid type."""
     router = SwarmRouter(swarm_type="SequentialWorkflow")
 
@@ -749,13 +794,12 @@ class TestSwarmRouterErrorHandling:
         router._create_swarm("Test task")
 
 
-class TestSwarmRouterCaching:
-    """Test swarm caching functionality."""
-
-    def test_swarm_caching(self):
+def test_swarm_caching():
     """Test that swarms are cached for performance."""
     agents = [Agent(agent_name="test", agent_description="test")]
-        router = SwarmRouter(agents=agents, swarm_type="SequentialWorkflow")
+    router = SwarmRouter(
+        agents=agents, swarm_type="SequentialWorkflow"
+    )
 
     # Create swarm first time
     swarm1 = router._create_swarm("Task 1")
@@ -766,10 +810,13 @@ class TestSwarmRouterCaching:
     # Should be the same cached instance
     assert swarm1 is swarm2
 
-    def test_swarm_cache_different_parameters(self):
+
+def test_swarm_cache_different_parameters():
     """Test that different parameters create different cached swarms."""
     agents = [Agent(agent_name="test", agent_description="test")]
-        router = SwarmRouter(agents=agents, swarm_type="SequentialWorkflow")
+    router = SwarmRouter(
+        agents=agents, swarm_type="SequentialWorkflow"
+    )
 
     # Create swarms with different parameters
     swarm1 = router._create_swarm("Task 1", param1="value1")
@@ -779,28 +826,26 @@ class TestSwarmRouterCaching:
     assert swarm1 is not swarm2
 
 
-class TestSwarmRouterOutputTypes:
-    """Test different output types."""
-
-    @pytest.fixture
-    def sample_agents(self):
-        """Create sample agents for testing."""
-        return [
-            Agent(
-                agent_name="TestAgent",
-                agent_description="Test agent",
-                system_prompt="You are a test agent.",
-                max_loops=1,
-            ),
-        ]
-
-    @pytest.mark.parametrize("output_type", [
-        "string", "str", "list", "json", "dict", "yaml", "xml",
-        "dict-all-except-first", "dict-first", "list-all-except-first"
-    ])
-    def test_output_types(self, sample_agents, output_type):
+@pytest.mark.parametrize(
+    "output_type",
+    [
+        "string",
+        "str",
+        "list",
+        "json",
+        "dict",
+        "yaml",
+        "xml",
+        "dict-all-except-first",
+        "dict-first",
+        "list-all-except-first",
+    ],
+)
+def test_output_types(sample_agents, output_type):
     """Test different output types."""
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = f"Response for {output_type}"
         mock_workflow.return_value = mock_instance
@@ -816,10 +861,7 @@ class TestSwarmRouterOutputTypes:
         assert result == f"Response for {output_type}"
 
 
-class TestSwarmRouterIntegration:
-    """Integration tests for SwarmRouter."""
-
-    def test_full_workflow_with_sequential_workflow(self):
+def test_full_workflow_with_sequential_workflow():
     """Test complete workflow with SequentialWorkflow."""
     agents = [
         Agent(
@@ -836,7 +878,9 @@ class TestSwarmRouterIntegration:
         ),
     ]
 
-        with patch('swarms.structs.sequential_workflow.SequentialWorkflow') as mock_workflow:
+    with patch(
+        "swarms.structs.sequential_workflow.SequentialWorkflow"
+    ) as mock_workflow:
         mock_instance = Mock()
         mock_instance.run.return_value = "Complete workflow response"
         mock_workflow.return_value = mock_instance
@@ -866,7 +910,8 @@ class TestSwarmRouterIntegration:
         assert isinstance(config_dict, dict)
         assert config_dict["swarm_type"] == "SequentialWorkflow"
 
-    def test_swarm_router_config_model(self):
+
+def test_swarm_router_config_model():
     """Test SwarmRouterConfig model."""
     config = SwarmRouterConfig(
         name="test-config",
@@ -882,12 +927,10 @@ class TestSwarmRouterIntegration:
     assert config.task == "Test task"
     assert config.multi_agent_collab_prompt is True
 
-    def test_document_model(self):
+
+def test_document_model():
     """Test Document model."""
-        doc = Document(
-            file_path="test.txt",
-            data="Test document content"
-        )
+    doc = Document(file_path="test.txt", data="Test document content")
 
     assert doc.file_path == "test.txt"
     assert doc.data == "Test document content"

tests/test_main_features.py

@@ -23,11 +23,8 @@ from swarms.structs import (
 from swarms.structs.hiearchical_swarm import HierarchicalSwarm
 from swarms.structs.tree_swarm import ForestSwarm, Tree, TreeAgent
 
-# Load environment variables
 load_dotenv()
 
-API_KEY = os.getenv("OPENAI_API_KEY")
-
 
 def generate_timestamp() -> str:
     """Generate a timestamp string for filenames"""
@@ -82,58 +79,6 @@ def write_markdown_report(
             f.write("---\n\n")
 
 
-# def create_github_issue(test_result: Dict[str, Any]) -> Dict[str, Any]:
-#     """Create a GitHub issue for a failed test"""
-#     if not all([GITHUB_TOKEN, GITHUB_REPO_OWNER, GITHUB_REPO_NAME]):
-#         logger.warning("GitHub credentials not configured. Skipping issue creation.")
-#         return None
-
-#     if test_result["status"] != "failed":
-#         return None
-
-#     issue_title = f"Automated Test Failure: {test_result['test_name']}"
-
-#     issue_body = f"""
-# ## Test Failure Report
-
-# - **Test Name**: `{test_result['test_name']}`
-# - **Timestamp**: `{datetime.now().isoformat()}`
-# - **Status**: {test_result['status']}
-
-# ### Error Information
-# ```
-# {test_result.get('error', 'No error message available')}
-# ```
-
-# ### Response (if available)
-# ```json
-# {json.dumps(test_result.get('response', {}), indent=2)}
-# ```
-
-# ---
-# *This issue was automatically generated by the Swarms testing workflow.*
-# """
-
-#     payload = {
-#         "title": issue_title,
-#         "body": issue_body,
-#         "labels": ["bug", "test-failure", "automated-report"],
-#     }
-
-#     try:
-#         response = requests.post(
-#             f"{BASE_URL}/repos/{GITHUB_REPO_OWNER}/{GITHUB_REPO_NAME}/issues",
-#             headers=GITHUB_HEADERS,
-#             json=payload,
-#         )
-#         response.raise_for_status()
-#         logger.info(f"Created GitHub issue for {test_result['test_name']}")
-#         return response.json()
-#     except requests.exceptions.RequestException as e:
-#         logger.error(f"Failed to create GitHub issue: {e.response.text if e.response else str(e)}")
-#         return None
-
-
 def create_test_agent(
     name: str,
     system_prompt: str = None,