From 9829d770d020fa6b41f3a023a5ba7fe3af2bdea7 Mon Sep 17 00:00:00 2001 From: Steve-Dusty Date: Wed, 29 Oct 2025 21:12:26 -0700 Subject: [PATCH 1/2] align tests with actual SwarmRouter API for test_swarm_router --- tests/structs/test_swarm_router.py | 970 ++++++++--------------------- 1 file changed, 276 insertions(+), 694 deletions(-) diff --git a/tests/structs/test_swarm_router.py b/tests/structs/test_swarm_router.py index 098ce3d5..3c31ec1d 100644 --- a/tests/structs/test_swarm_router.py +++ b/tests/structs/test_swarm_router.py @@ -1,34 +1,46 @@ -from unittest.mock import Mock, patch - import pytest -from swarms.structs.agent import Agent from swarms.structs.swarm_router import ( SwarmRouter, SwarmRouterConfig, - SwarmRouterConfigError, SwarmRouterRunError, + SwarmRouterConfigError, + Document, ) +from swarms.structs.agent import Agent + +# ============================================================================ +# Helper Functions +# ============================================================================ -@pytest.fixture -def sample_agents(): +def create_sample_agents(): """Create sample agents for testing.""" return [ Agent( agent_name="ResearchAgent", agent_description="Specializes in researching topics", system_prompt="You are a research specialist.", + model_name="gpt-4o-mini", max_loops=1, + verbose=False, + print_on=False, ), Agent( agent_name="CodeAgent", agent_description="Expert in coding", system_prompt="You are a coding expert.", + model_name="gpt-4o-mini", max_loops=1, + verbose=False, + print_on=False, ), ] +# ============================================================================ +# Initialization Tests +# ============================================================================ + def test_default_initialization(): """Test SwarmRouter with default parameters.""" @@ -52,8 +64,10 @@ def test_default_initialization(): assert router.telemetry_enabled is False -def test_custom_initialization(sample_agents): +def test_custom_initialization(): """Test SwarmRouter with custom parameters.""" + sample_agents = create_sample_agents() + router = SwarmRouter( name="test-router", description="Test router description", @@ -66,8 +80,8 @@ def test_custom_initialization(sample_agents): rules="Test rules", documents=["doc1.txt", "doc2.txt"], output_type="json", - verbose=True, - telemetry_enabled=True, + verbose=False, # Keep quiet for tests + telemetry_enabled=False, ) assert router.name == "test-router" @@ -81,815 +95,383 @@ def test_custom_initialization(sample_agents): assert router.rules == "Test rules" assert router.documents == ["doc1.txt", "doc2.txt"] assert router.output_type == "json" - assert router.verbose is True - assert router.telemetry_enabled is True + assert router.verbose is False + assert router.telemetry_enabled is False -def test_initialization_with_heavy_swarm_config(sample_agents): +def test_initialization_with_heavy_swarm_config(): """Test SwarmRouter with HeavySwarm specific configuration.""" + sample_agents = create_sample_agents() + router = SwarmRouter( agents=sample_agents, swarm_type="HeavySwarm", heavy_swarm_loops_per_agent=2, - heavy_swarm_question_agent_model_name="gpt-4", - heavy_swarm_worker_model_name="gpt-3.5-turbo", + heavy_swarm_question_agent_model_name="gpt-4o-mini", + heavy_swarm_worker_model_name="gpt-4o-mini", heavy_swarm_swarm_show_output=False, ) assert router.swarm_type == "HeavySwarm" assert router.heavy_swarm_loops_per_agent == 2 - assert router.heavy_swarm_question_agent_model_name == "gpt-4" - assert router.heavy_swarm_worker_model_name == "gpt-3.5-turbo" + assert router.heavy_swarm_question_agent_model_name == "gpt-4o-mini" + assert router.heavy_swarm_worker_model_name == "gpt-4o-mini" assert router.heavy_swarm_swarm_show_output is False -def test_initialization_with_council_judge_config(): - """Test SwarmRouter with CouncilAsAJudge specific configuration.""" - router = SwarmRouter( - swarm_type="CouncilAsAJudge", - council_judge_model_name="gpt-4o", - ) - - assert router.swarm_type == "CouncilAsAJudge" - assert router.council_judge_model_name == "gpt-4o" - +def test_initialization_with_agent_rearrange_config(): + """Test SwarmRouter with AgentRearrange specific configuration.""" + sample_agents = create_sample_agents() -def test_initialization_with_agent_rearrange_flow(sample_agents): - """Test SwarmRouter with AgentRearrange and flow configuration.""" - flow = "agent1 -> agent2 -> agent1" router = SwarmRouter( agents=sample_agents, swarm_type="AgentRearrange", - rearrange_flow=flow, + rearrange_flow="ResearchAgent -> CodeAgent", ) assert router.swarm_type == "AgentRearrange" - assert router.rearrange_flow == flow - - -def test_invalid_swarm_type(): - """Test error when invalid swarm type is provided.""" - with pytest.raises(ValueError): - SwarmRouter(swarm_type="InvalidSwarmType") + assert router.rearrange_flow == "ResearchAgent -> CodeAgent" +# ============================================================================ +# Configuration Tests +# ============================================================================ -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(): - """Test error when no rearrange_flow provided for AgentRearrange.""" - agents = [Agent(agent_name="test", agent_description="test")] - with pytest.raises(SwarmRouterConfigError): - SwarmRouter( - agents=agents, - swarm_type="AgentRearrange", - rearrange_flow=None, - ) - - -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(): - """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(): - """Test CouncilAsAJudge can be created without agents.""" - router = SwarmRouter(swarm_type="CouncilAsAJudge", agents=None) - assert router.swarm_type == "CouncilAsAJudge" - - -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", - ] - - for swarm_type in expected_types: - assert swarm_type in factory - assert callable(factory[swarm_type]) - -def test_create_heavy_swarm(sample_agents): - """Test HeavySwarm creation.""" +def test_initialization_with_shared_memory(): + """Test SwarmRouter with shared memory system.""" + sample_agents = create_sample_agents() + router = SwarmRouter( agents=sample_agents, - swarm_type="HeavySwarm", - heavy_swarm_loops_per_agent=2, - heavy_swarm_question_agent_model_name="gpt-4", - heavy_swarm_worker_model_name="gpt-3.5-turbo", + shared_memory_system=None, # Test with None for now ) - swarm = router._create_heavy_swarm() - assert swarm is not None - assert hasattr(swarm, "run") + assert router.shared_memory_system is None -def test_create_agent_rearrange(sample_agents): - """Test AgentRearrange creation.""" +def test_initialization_with_worker_tools(): + """Test SwarmRouter with worker tools.""" + sample_agents = create_sample_agents() + router = SwarmRouter( agents=sample_agents, - swarm_type="AgentRearrange", - rearrange_flow="agent1 -> agent2", + worker_tools=[], # Empty list for now ) - swarm = router._create_agent_rearrange() - assert swarm is not None - assert hasattr(swarm, "run") - + assert router.worker_tools == [] -def test_create_sequential_workflow(sample_agents): - """Test SequentialWorkflow creation.""" - router = SwarmRouter( - agents=sample_agents, swarm_type="SequentialWorkflow" - ) +# ============================================================================ +# Document Management Tests +# ============================================================================ - swarm = router._create_sequential_workflow() - assert swarm is not None - assert hasattr(swarm, "run") - -def test_create_concurrent_workflow(sample_agents): - """Test ConcurrentWorkflow creation.""" - router = SwarmRouter( - agents=sample_agents, swarm_type="ConcurrentWorkflow" +def test_document_creation(): + """Test Document creation.""" + doc = Document( + file_path="/path/to/test/document.txt", + data="This is test content", ) - swarm = router._create_concurrent_workflow() - assert swarm is not None - assert hasattr(swarm, "run") - - -def test_create_group_chat(sample_agents): - """Test GroupChat creation.""" - router = SwarmRouter(agents=sample_agents, swarm_type="GroupChat") + assert doc.file_path == "/path/to/test/document.txt" + assert doc.data == "This is test content" - swarm = router._create_group_chat() - assert swarm is not None - assert hasattr(swarm, "run") +def test_router_with_documents(): + """Test SwarmRouter with document configuration.""" + sample_agents = create_sample_agents() + documents = [ + Document(file_path="/path/to/doc1.txt", data="Content1"), + Document(file_path="/path/to/doc2.txt", data="Content2"), + ] -def test_create_multi_agent_router(sample_agents): - """Test MultiAgentRouter creation.""" router = SwarmRouter( - agents=sample_agents, swarm_type="MultiAgentRouter" + agents=sample_agents, + documents=documents, ) - swarm = router._create_multi_agent_router() - assert swarm is not None - assert hasattr(swarm, "run") - + assert len(router.documents) == 2 + assert router.documents[0].file_path == "/path/to/doc1.txt" + assert router.documents[1].file_path == "/path/to/doc2.txt" -def test_create_mixture_of_agents(sample_agents): - """Test MixtureOfAgents creation.""" - router = SwarmRouter( - agents=sample_agents, swarm_type="MixtureOfAgents" - ) +# ============================================================================ +# Configuration Class Tests +# ============================================================================ - swarm = router._create_mixture_of_agents() - assert swarm is not None - assert hasattr(swarm, "run") - -def test_create_majority_voting(sample_agents): - """Test MajorityVoting creation.""" - router = SwarmRouter( - agents=sample_agents, swarm_type="MajorityVoting" +def test_swarm_router_config_creation(): + """Test SwarmRouterConfig creation.""" + config = SwarmRouterConfig( + name="test-config", + description="Test configuration", + swarm_type="SequentialWorkflow", + rearrange_flow=None, + rules=None, + multi_agent_collab_prompt=True, + task="Test task", ) - swarm = router._create_majority_voting() - assert swarm is not None - assert hasattr(swarm, "run") - - -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") - - -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 config.name == "test-config" + assert config.description == "Test configuration" + assert config.swarm_type == "SequentialWorkflow" + assert config.task == "Test task" -def test_create_interactive_group_chat(sample_agents): - """Test InteractiveGroupChat creation.""" - router = SwarmRouter( - agents=sample_agents, swarm_type="InteractiveGroupChat" +def test_router_with_config(): + """Test SwarmRouter initialization matches config structure.""" + sample_agents = create_sample_agents() + config = SwarmRouterConfig( + name="config-router", + description="Router from config", + swarm_type="SequentialWorkflow", + rearrange_flow=None, + rules="Test rules", + multi_agent_collab_prompt=False, + task="Test task", ) - swarm = router._create_interactive_group_chat() - assert swarm is not None - assert hasattr(swarm, "run") - + # SwarmRouter doesn't accept config directly, but we can verify config is valid + assert config.name == "config-router" + assert config.description == "Router from config" + assert config.swarm_type == "SequentialWorkflow" + assert config.rules == "Test rules" -def test_create_hierarchical_swarm(sample_agents): - """Test HierarchicalSwarm creation.""" + # Create router with matching parameters router = SwarmRouter( - agents=sample_agents, swarm_type="HiearchicalSwarm" + name=config.name, + description=config.description, + agents=sample_agents, + swarm_type=config.swarm_type, + rules=config.rules, ) - swarm = router._create_hierarchical_swarm() - assert swarm is not None - assert hasattr(swarm, "run") + assert router.name == config.name + assert router.description == config.description + assert router.swarm_type == config.swarm_type + assert router.rules == config.rules + +# ============================================================================ +# Basic Execution Tests +# ============================================================================ -def test_create_batched_grid_workflow(sample_agents): - """Test BatchedGridWorkflow creation.""" +def test_run_with_sequential_workflow(): + """Test running SwarmRouter with SequentialWorkflow.""" + sample_agents = create_sample_agents() + router = SwarmRouter( - agents=sample_agents, swarm_type="BatchedGridWorkflow" + agents=sample_agents, + swarm_type="SequentialWorkflow", + verbose=False, ) - swarm = router._create_batched_grid_workflow() - assert swarm is not None - assert hasattr(swarm, "run") - - -@pytest.mark.parametrize( - "swarm_type", - [ - "SequentialWorkflow", - "ConcurrentWorkflow", - "GroupChat", - "MultiAgentRouter", - "MixtureOfAgents", - "MajorityVoting", - "MALT", - "CouncilAsAJudge", - "InteractiveGroupChat", - "HiearchicalSwarm", - "BatchedGridWorkflow", - ], -) -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: - # Mock the LLM to return a simple response - mock_llm_instance = Mock() - mock_llm_instance.agenerate.return_value = ( - "Test response from agent" - ) - mock_llm.return_value = mock_llm_instance - - router = SwarmRouter( - agents=sample_agents, - swarm_type=swarm_type, - max_loops=1, - ) - - # Test with a simple task - task = "Write a simple Python function to add two numbers" - - try: - result = router.run(task) - 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 - router = SwarmRouter( - agents=sample_agents, - swarm_type=swarm_type, - rearrange_flow="agent1 -> agent2", - max_loops=1, - ) - result = router.run(task) - assert ( - result is not None - ), f"Swarm type {swarm_type} returned None result" - - -def test_heavy_swarm_execution(): - """Test HeavySwarm execution.""" - 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 - - router = SwarmRouter(swarm_type="HeavySwarm") - - result = router.run("Test task") - assert result is not None - assert result == "HeavySwarm response" - - -def test_agent_rearrange_execution(sample_agents): - """Test AgentRearrange execution with flow.""" - 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 - - router = SwarmRouter( - agents=sample_agents, - swarm_type="AgentRearrange", - rearrange_flow="agent1 -> agent2", - ) - - result = router.run("Test task") - assert result is not None - assert result == "AgentRearrange response" - - -def test_run_method(sample_agents): - """Test basic run method.""" - 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" - ) - - result = router.run("Test task") - assert result is not None - assert result == "Test response" - - -def test_run_with_image(sample_agents): - """Test run method with image input.""" - 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" - ) - - 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(sample_agents): - """Test run method with tasks list.""" - 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" - ) - - 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(sample_agents): - """Test batch_run method.""" - 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" - ) - - tasks = ["Task 1", "Task 2", "Task 3"] - results = router.batch_run(tasks) - - assert results is not None - assert len(results) == 3 - assert all(result is not None for result in results) + result = router.run("What is 2+2?") + assert result is not None -def test_concurrent_run_method(sample_agents): - """Test concurrent_run method.""" - 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" - ) +def test_run_with_no_agents(): + """Test running SwarmRouter with no agents.""" + router = SwarmRouter() - result = router.concurrent_run("Test task") - assert result is not None - assert result == "Concurrent response" + with pytest.raises(RuntimeError): + router.run("Test task") -def test_call_method(sample_agents): - """Test __call__ method.""" - 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 +def test_run_with_empty_task(): + """Test running SwarmRouter with empty task.""" + sample_agents = create_sample_agents() - router = SwarmRouter( - agents=sample_agents, swarm_type="SequentialWorkflow" - ) + router = SwarmRouter(agents=sample_agents, verbose=False) - result = router("Test task") - assert result is not None - assert result == "Call response" + # Empty task is allowed, router will pass it to the swarm + result = router.run("") + assert result is not None -def test_call_with_image(sample_agents): - """Test __call__ method with image.""" - 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 +def test_run_with_none_task(): + """Test running SwarmRouter with None task.""" + sample_agents = create_sample_agents() - router = SwarmRouter( - agents=sample_agents, swarm_type="SequentialWorkflow" - ) + router = SwarmRouter(agents=sample_agents, verbose=False) - result = router("Test task", img="test.jpg") - assert result is not None - assert result == "Image call response" + # None task is allowed, router will pass it to the swarm + result = router.run(None) + assert result is not None +# ============================================================================ +# Batch Processing Tests +# ============================================================================ -def test_call_with_images_list(sample_agents): - """Test __call__ method with images list.""" - 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" - ) - - result = router("Test task", imgs=["test1.jpg", "test2.jpg"]) - assert result is not None - assert result == "Images call response" - - -def test_to_dict_method(sample_agents): - """Test to_dict method serialization.""" +def test_batch_run_with_tasks(): + """Test batch processing with multiple tasks.""" + sample_agents = create_sample_agents() + router = SwarmRouter( agents=sample_agents, - name="test-router", - description="Test description", - swarm_type="SequentialWorkflow", + verbose=False, ) - result_dict = router.to_dict() + tasks = ["What is 1+1?", "What is 2+2?"] + results = router.batch_run(tasks) + + assert len(results) == 2 + assert all(result is not None for result in results) - assert isinstance(result_dict, dict) - assert result_dict["name"] == "test-router" - assert result_dict["description"] == "Test description" - assert result_dict["swarm_type"] == "SequentialWorkflow" - assert "agents" in result_dict +def test_batch_run_with_empty_tasks(): + """Test batch processing with empty task list.""" + sample_agents = create_sample_agents() + + router = SwarmRouter(agents=sample_agents) + + results = router.batch_run([]) + assert results == [] -def test_activate_ape(sample_agents): - """Test activate_ape method.""" - router = SwarmRouter( - agents=sample_agents, - auto_generate_prompts=True, - ) - # Mock the auto_generate_prompt attribute - for agent in router.agents: - agent.auto_generate_prompt = False +def test_batch_run_with_no_agents(): + """Test batch processing with no agents.""" + router = SwarmRouter() - router.activate_ape() + with pytest.raises(RuntimeError): + router.batch_run(["Test task"]) - # Check that APE was activated for agents that support it - for agent in router.agents: - if hasattr(agent, "auto_generate_prompt"): - assert agent.auto_generate_prompt is True +# ============================================================================ +# Call Method Tests +# ============================================================================ -def test_handle_rules(sample_agents): - """Test handle_rules method.""" - rules = "Always be helpful and accurate." +def test_call_method(): + """Test __call__ method.""" + sample_agents = create_sample_agents() + router = SwarmRouter( agents=sample_agents, - rules=rules, + verbose=False, ) - 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}" - ) + result = router("What is the capital of France?") + assert result is not None -def test_update_system_prompt_for_agent_in_swarm(sample_agents): - """Test update_system_prompt_for_agent_in_swarm method.""" +def test_call_with_image(): + """Test __call__ method with image.""" + sample_agents = create_sample_agents() + router = SwarmRouter( agents=sample_agents, - multi_agent_collab_prompt=True, - ) - - 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(sample_agents): - """Test agent_config method.""" - router = SwarmRouter(agents=sample_agents) - - config = router.agent_config() - - assert isinstance(config, dict) - assert len(config) == len(sample_agents) - - for agent in sample_agents: - assert agent.agent_name in config - - -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" + verbose=False, ) - router.swarm = Mock() - router.swarm.conversation = mock_conversation - result = router.fetch_message_history_as_string() - assert result == "Test history" + # Test with None image (no actual image processing) + result = router("Describe this image", img=None) + assert result is not None +# ============================================================================ +# Output Type Tests +# ============================================================================ -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" - ) - - result = router.fetch_message_history_as_string() - assert result is None - - -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(): - """Test error handling during task execution.""" - agents = [Agent(agent_name="test", agent_description="test")] - - with patch( - "swarms.structs.sequential_workflow.SequentialWorkflow" - ) as mock_workflow: - mock_instance = Mock() - mock_instance.run.side_effect = Exception( - "Test execution error" - ) - mock_workflow.return_value = mock_instance +def test_different_output_types(): + """Test router with different output types.""" + sample_agents = create_sample_agents() + + for output_type in ["dict", "json", "string", "list"]: router = SwarmRouter( - agents=agents, swarm_type="SequentialWorkflow" + agents=sample_agents, + output_type=output_type, + verbose=False, ) + + result = router.run("Simple test task") + assert result is not None - with pytest.raises(SwarmRouterRunError): - router.run("Test task") - - -def test_batch_run_error_handling(): - """Test error handling during batch execution.""" - agents = [Agent(agent_name="test", agent_description="test")] +# ============================================================================ +# Error Handling Tests +# ============================================================================ - 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" - ) +def test_swarm_router_run_error(): + """Test SwarmRouterRunError exception.""" + error = SwarmRouterRunError("Test error message") + assert str(error) == "Test error message" - with pytest.raises(RuntimeError): - router.batch_run(["Task 1", "Task 2"]) +def test_swarm_router_config_error(): + """Test SwarmRouterConfigError exception.""" + error = SwarmRouterConfigError("Config error message") + assert str(error) == "Config error message" -def test_invalid_swarm_type_error(): - """Test error when creating swarm with invalid type.""" - router = SwarmRouter(swarm_type="SequentialWorkflow") - - # Manually set an invalid swarm type to test the factory - router.swarm_type = "InvalidType" - - with pytest.raises(ValueError): - router._create_swarm("Test task") +def test_invalid_swarm_type(): + """Test router with invalid swarm type.""" + sample_agents = create_sample_agents() -def test_swarm_caching(): - """Test that swarms are cached for performance.""" - agents = [Agent(agent_name="test", agent_description="test")] + # This should not raise an error during initialization router = SwarmRouter( - agents=agents, swarm_type="SequentialWorkflow" + agents=sample_agents, + swarm_type="InvalidSwarmType", ) - # Create swarm first time - swarm1 = router._create_swarm("Task 1") - - # Create swarm second time with same parameters - swarm2 = router._create_swarm("Task 1") + # But should raise ValueError during execution when creating swarm + with pytest.raises(ValueError, match="Invalid swarm type: InvalidSwarmType"): + router.run("Test task") - # Should be the same cached instance - assert swarm1 is swarm2 +# ============================================================================ +# Integration Tests +# ============================================================================ -def test_swarm_cache_different_parameters(): - """Test that different parameters create different cached swarms.""" - agents = [Agent(agent_name="test", agent_description="test")] +def test_complete_workflow(): + """Test complete workflow from initialization to execution.""" + # Create agents + agents = create_sample_agents() + + # Create router with configuration router = SwarmRouter( - agents=agents, swarm_type="SequentialWorkflow" + name="integration-test-router", + description="Router for integration testing", + agents=agents, + swarm_type="SequentialWorkflow", + max_loops=1, + verbose=False, + output_type="string", ) - - # Create swarms with different parameters - swarm1 = router._create_swarm("Task 1", param1="value1") - swarm2 = router._create_swarm("Task 2", param1="value2") - - # Should be different instances - assert swarm1 is not swarm2 - - -@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: - mock_instance = Mock() - mock_instance.run.return_value = f"Response for {output_type}" - mock_workflow.return_value = mock_instance - - router = SwarmRouter( - agents=sample_agents, - swarm_type="SequentialWorkflow", - output_type=output_type, - ) - - result = router.run("Test task") - assert result is not None - assert result == f"Response for {output_type}" - - -def test_full_workflow_with_sequential_workflow(): - """Test complete workflow with SequentialWorkflow.""" - agents = [ - Agent( - agent_name="ResearchAgent", - agent_description="Research specialist", - system_prompt="You are a research specialist.", - max_loops=1, - ), - Agent( - agent_name="CodeAgent", - agent_description="Code expert", - system_prompt="You are a code expert.", - max_loops=1, - ), + + # Execute single task + result = router.run("Calculate the sum of 5 and 7") + assert result is not None + + # Execute batch tasks + tasks = [ + "What is 10 + 15?", + "What is 20 - 8?", + "What is 6 * 7?", ] + batch_results = router.batch_run(tasks) + + assert len(batch_results) == 3 + assert all(result is not None for result in batch_results) - 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 - - router = SwarmRouter( - agents=agents, - swarm_type="SequentialWorkflow", - max_loops=2, - rules="Always provide accurate information", - multi_agent_collab_prompt=True, - verbose=True, - ) - - # Test various methods - result = router.run("Research and code a Python function") - assert result is not None - - batch_results = router.batch_run(["Task 1", "Task 2"]) - assert batch_results is not None - assert len(batch_results) == 2 - - concurrent_result = router.concurrent_run("Concurrent task") - assert concurrent_result is not None - - # Test serialization - config_dict = router.to_dict() - assert isinstance(config_dict, dict) - assert config_dict["swarm_type"] == "SequentialWorkflow" - - -def test_swarm_router_config_model(): - """Test SwarmRouterConfig model.""" - config = SwarmRouterConfig( - name="test-config", - description="Test configuration", - swarm_type="SequentialWorkflow", - task="Test task", - multi_agent_collab_prompt=True, - ) - assert config.name == "test-config" - assert config.description == "Test configuration" - assert config.swarm_type == "SequentialWorkflow" - assert config.task == "Test task" - assert config.multi_agent_collab_prompt is True +def test_router_reconfiguration(): + """Test reconfiguring router after initialization.""" + sample_agents = create_sample_agents() + + router = SwarmRouter(agents=sample_agents) + + # Change configuration + router.max_loops = 3 + router.output_type = "json" + router.verbose = False + + assert router.max_loops == 3 + assert router.output_type == "json" + assert router.verbose is False + + # Test execution with new configuration + result = router.run("Test reconfiguration") + assert result is not None if __name__ == "__main__": - pytest.main([__file__, "-v"]) + pytest.main([__file__, "-v"]) \ No newline at end of file From e5e5248d293edf7fae0381d60484e69b2c07cbae Mon Sep 17 00:00:00 2001 From: Steve-Dusty Date: Wed, 29 Oct 2025 23:03:28 -0700 Subject: [PATCH 2/2] refactored test_agent_rearrange --- tests/structs/test_agent_rearrange.py | 547 ++++++++++++++++++++++---- 1 file changed, 472 insertions(+), 75 deletions(-) diff --git a/tests/structs/test_agent_rearrange.py b/tests/structs/test_agent_rearrange.py index 7a9b256c..ffb53278 100644 --- a/tests/structs/test_agent_rearrange.py +++ b/tests/structs/test_agent_rearrange.py @@ -1,128 +1,525 @@ import pytest -from unittest.mock import MagicMock -from swarms import AgentRearrange +from swarms import Agent, AgentRearrange -class MockAgent: - def __init__(self, name): - self.name = name +# ============================================================================ +# Helper Functions +# ============================================================================ - def run(self, task, img=None, *args, **kwargs): - return f"{self.name} processed {task}" - -@pytest.fixture -def mock_agents(): +def create_sample_agents(): + """Create sample agents for testing.""" return [ - MockAgent(name="Agent1"), - MockAgent(name="Agent2"), - MockAgent(name="Agent3"), + Agent( + agent_name="ResearchAgent", + agent_description="Specializes in researching topics", + system_prompt="You are a research specialist. Provide concise answers.", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), + Agent( + agent_name="WriterAgent", + agent_description="Expert in writing content", + system_prompt="You are a writing expert. Provide concise answers.", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), + Agent( + agent_name="ReviewerAgent", + agent_description="Expert in reviewing content", + system_prompt="You are a review expert. Provide concise answers.", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), ] -@pytest.fixture -def agent_rearrange(mock_agents): - return AgentRearrange( - agents=mock_agents, flow="Agent1 -> Agent2 -> Agent3" - ) +# ============================================================================ +# Initialization Tests +# ============================================================================ -def test_initialization(mock_agents): +def test_initialization(): + """Test AgentRearrange initialization.""" + agents = create_sample_agents() + agent_rearrange = AgentRearrange( - agents=mock_agents, flow="Agent1 -> Agent2 -> Agent3" + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + max_loops=1, + verbose=True, ) + assert len(agent_rearrange.agents) == 3 - assert agent_rearrange.flow == "Agent1 -> Agent2 -> Agent3" + assert agent_rearrange.flow == "ResearchAgent -> WriterAgent -> ReviewerAgent" + assert agent_rearrange.max_loops == 1 + assert agent_rearrange.verbose is True + print("✓ test_initialization passed") + + +def test_initialization_with_team_awareness(): + """Test AgentRearrange with team_awareness enabled.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + team_awareness=True, + verbose=True, + ) + + assert agent_rearrange.flow == "ResearchAgent -> WriterAgent -> ReviewerAgent" + print("✓ test_initialization_with_team_awareness passed") + + +def test_initialization_with_custom_output_type(): + """Test AgentRearrange with custom output types.""" + agents = create_sample_agents() + + for output_type in ["all", "final", "list", "dict"]: + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + output_type=output_type, + verbose=True, + ) + assert agent_rearrange.output_type == output_type + + print("✓ test_initialization_with_custom_output_type passed") + + +# ============================================================================ +# Agent Management Tests +# ============================================================================ -def test_add_agent(agent_rearrange): - new_agent = MockAgent(name="Agent4") +def test_add_agent(): + """Test adding an agent to AgentRearrange.""" + agents = create_sample_agents()[:2] # Only 2 agents + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + verbose=True, + ) + + new_agent = Agent( + agent_name="EditorAgent", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ) + agent_rearrange.add_agent(new_agent) - assert "Agent4" in agent_rearrange.agents + assert "EditorAgent" in agent_rearrange.agents + assert len(agent_rearrange.agents) == 3 + print("✓ test_add_agent passed") -def test_remove_agent(agent_rearrange): - agent_rearrange.remove_agent("Agent2") - assert "Agent2" not in agent_rearrange.agents +def test_remove_agent(): + """Test removing an agent from AgentRearrange.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, + ) + + agent_rearrange.remove_agent("ReviewerAgent") + assert "ReviewerAgent" not in agent_rearrange.agents + assert len(agent_rearrange.agents) == 2 + print("✓ test_remove_agent passed") -def test_add_agents(agent_rearrange): - new_agents = [MockAgent(name="Agent4"), MockAgent(name="Agent5")] +def test_add_agents(): + """Test adding multiple agents to AgentRearrange.""" + agents = create_sample_agents()[:1] # Start with 1 agent + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent", + verbose=True, + ) + + new_agents = [ + Agent( + agent_name="Agent4", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), + Agent( + agent_name="Agent5", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), + ] + agent_rearrange.add_agents(new_agents) assert "Agent4" in agent_rearrange.agents assert "Agent5" in agent_rearrange.agents + assert len(agent_rearrange.agents) == 3 + print("✓ test_add_agents passed") -def test_validate_flow_valid(agent_rearrange): +# ============================================================================ +# Flow Validation Tests +# ============================================================================ + + +def test_validate_flow_valid(): + """Test flow validation with valid flow.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, + ) + assert agent_rearrange.validate_flow() is True + print("✓ test_validate_flow_valid passed") + + +def test_validate_flow_invalid(): + """Test flow validation with invalid agent name in flow.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, + ) + # Change to invalid flow + agent_rearrange.flow = "ResearchAgent -> NonExistentAgent" -def test_validate_flow_invalid(agent_rearrange): - agent_rearrange.flow = "Agent1 -> Agent4" - with pytest.raises(ValueError): + try: agent_rearrange.validate_flow() + assert False, "Should have raised ValueError" + except ValueError as e: + assert "not registered" in str(e) + print("✓ test_validate_flow_invalid passed") + + +def test_validate_flow_no_arrow(): + """Test flow validation without arrow syntax.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + verbose=True, + ) + + agent_rearrange.flow = "ResearchAgent WriterAgent" + + try: + agent_rearrange.validate_flow() + assert False, "Should have raised ValueError" + except ValueError as e: + assert "'->" in str(e) + print("✓ test_validate_flow_no_arrow passed") + + +# ============================================================================ +# Flow Pattern Tests +# ============================================================================ + + +def test_set_custom_flow(): + """Test setting custom flow.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + verbose=True, + ) + new_flow = "WriterAgent -> ResearchAgent -> ReviewerAgent" + agent_rearrange.set_custom_flow(new_flow) + assert agent_rearrange.flow == new_flow + print("✓ test_set_custom_flow passed") -def test_run(agent_rearrange): - result = agent_rearrange.run("Test Task") - assert ( - result - == "Agent1 processed Test Task; Agent2 processed Agent1 processed Test Task; Agent3 processed Agent2 processed Agent1 processed Test Task" + +def test_sequential_flow(): + """Test sequential flow execution.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, ) + result = agent_rearrange.run("What is 2+2?") + assert result is not None + print("✓ test_sequential_flow passed") + -def test_run_with_custom_tasks(agent_rearrange): - custom_tasks = {"Agent2": "Custom Task"} - result = agent_rearrange.run( - "Test Task", custom_tasks=custom_tasks +def test_concurrent_flow(): + """Test concurrent flow execution with comma syntax.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent, WriterAgent -> ReviewerAgent", + max_loops=1, + verbose=True, ) - assert ( - result - == "Agent1 processed Test Task; Agent2 processed Custom Task; Agent3 processed Agent2 processed Custom Task" + + result = agent_rearrange.run("What is 3+3?") + assert result is not None + print("✓ test_concurrent_flow passed") + + +# ============================================================================ +# Sequential Awareness Tests +# ============================================================================ + + +def test_get_sequential_flow_structure(): + """Test getting sequential flow structure.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, ) + flow_structure = agent_rearrange.get_sequential_flow_structure() + assert flow_structure is not None + assert isinstance(flow_structure, str) + print("✓ test_get_sequential_flow_structure passed") + -def test_run_with_human_intervention(agent_rearrange): - agent_rearrange.human_in_the_loop = True - agent_rearrange.custom_human_in_the_loop = MagicMock( - return_value="Human processed Task" +def test_get_agent_sequential_awareness(): + """Test getting agent sequential awareness.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, ) - agent_rearrange.flow = "Agent1 -> H -> Agent3" - result = agent_rearrange.run("Test Task") - assert ( - result - == "Agent1 processed Test Task; Human processed Task; Agent3 processed Human processed Task" + + awareness = agent_rearrange.get_agent_sequential_awareness("WriterAgent") + assert awareness is not None + assert isinstance(awareness, str) + print("✓ test_get_agent_sequential_awareness passed") + + +# ============================================================================ +# Execution Tests +# ============================================================================ + + +def test_run_basic(): + """Test basic run execution.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, ) + result = agent_rearrange.run("Calculate 5+5") + assert result is not None + print("✓ test_run_basic passed") + -def test_run_sub_swarm(agent_rearrange): - sub_swarm_flow = "Agent1 -> Agent3" - agent_rearrange.add_sub_swarm("SubSwarm1", sub_swarm_flow) - result = agent_rearrange.run_sub_swarm( - "SubSwarm1", "Sub Task", None +def test_run_with_different_output_types(): + """Test run with different output types.""" + agents = create_sample_agents() + + for output_type in ["all", "final"]: + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + output_type=output_type, + max_loops=1, + verbose=True, + ) + + result = agent_rearrange.run("What is the capital of France?") + assert result is not None + + print("✓ test_run_with_different_output_types passed") + + +def test_callable_execution(): + """Test __call__ method.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, ) - assert ( - result - == "Agent1 processed Sub Task; Agent3 processed Agent1 processed Sub Task" + + result = agent_rearrange("What is 10+10?") + assert result is not None + print("✓ test_callable_execution passed") + + +# ============================================================================ +# Batch Processing Tests +# ============================================================================ + + +def test_batch_run(): + """Test batch processing.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, ) + tasks = ["What is 1+1?", "What is 2+2?"] + results = agent_rearrange.batch_run(tasks, batch_size=2) + + assert results is not None + assert len(results) == 2 + print("✓ test_batch_run passed") + + +def test_concurrent_run(): + """Test concurrent execution.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, + ) + + tasks = ["What is 3+3?", "What is 4+4?"] + results = agent_rearrange.concurrent_run(tasks, max_workers=2) + + assert results is not None + assert len(results) == 2 + print("✓ test_concurrent_run passed") -def test_process_agent_or_swarm(agent_rearrange): - result = agent_rearrange.process_agent_or_swarm( - "Agent1", "Process Task", None + +# ============================================================================ +# Serialization Tests +# ============================================================================ + + +def test_to_dict(): + """Test serialization to dictionary.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + verbose=True, ) - assert result == "Agent1 processed Process Task" + + result_dict = agent_rearrange.to_dict() + assert isinstance(result_dict, dict) + assert "name" in result_dict + assert "flow" in result_dict + assert "agents" in result_dict + print("✓ test_to_dict passed") -def test_track_history(agent_rearrange): - agent_rearrange.track_history("Agent1", "Task Result") - assert agent_rearrange.swarm_history["Agent1"] == ["Task Result"] +# ============================================================================ +# Integration Tests +# ============================================================================ -def test_human_intervention(agent_rearrange): - agent_rearrange.human_in_the_loop = True - agent_rearrange.custom_human_in_the_loop = MagicMock( - return_value="Human processed Task" +def test_complete_workflow(): + """Test complete workflow with all features.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + name="test-workflow", + description="Test complete workflow", + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + max_loops=1, + team_awareness=True, + verbose=True, + output_type="all", ) - result = agent_rearrange.human_intervention("Task") - assert result == "Human processed Task" + + # Single run + result1 = agent_rearrange.run("What is Python?") + assert result1 is not None + + # Get flow structure + flow_structure = agent_rearrange.get_sequential_flow_structure() + assert flow_structure is not None + + # Serialize + result_dict = agent_rearrange.to_dict() + assert isinstance(result_dict, dict) + + print("✓ test_complete_workflow passed") + + +def main(): + """Run all tests.""" + tests = [ + test_initialization, + test_initialization_with_team_awareness, + test_initialization_with_custom_output_type, + test_add_agent, + test_remove_agent, + test_add_agents, + test_validate_flow_valid, + test_validate_flow_invalid, + test_validate_flow_no_arrow, + test_set_custom_flow, + test_sequential_flow, + test_concurrent_flow, + test_get_sequential_flow_structure, + test_get_agent_sequential_awareness, + test_run_basic, + test_run_with_different_output_types, + test_callable_execution, + test_batch_run, + test_concurrent_run, + test_to_dict, + test_complete_workflow, + ] + + print("="*60) + print("Running AgentRearrange Tests") + print("="*60) + + for test in tests: + try: + test() + except Exception as e: + print(f"✗ {test.__name__} failed: {e}") + + print("="*60) + print("All tests completed!") + print("="*60) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])