refactored test_agent_rearrange

tests/structs/test_agent_rearrange.py

@@ -1,128 +1,525 @@
 import pytest
-from unittest.mock import MagicMock
+from swarms import Agent, AgentRearrange
-from swarms import AgentRearrange
 
 
-class MockAgent:
+# ============================================================================
-    def __init__(self, name):
+# Helper Functions
-        self.name = name
+# ============================================================================
 
-    def run(self, task, img=None, *args, **kwargs):
-        return f"{self.name} processed {task}"
 
-
+def create_sample_agents():
-@pytest.fixture
+    """Create sample agents for testing."""
-def mock_agents():
     return [
-        MockAgent(name="Agent1"),
+        Agent(
-        MockAgent(name="Agent2"),
+            agent_name="ResearchAgent",
-        MockAgent(name="Agent3"),
+            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):
+# Initialization Tests
-    return AgentRearrange(
+# ============================================================================
-        agents=mock_agents, flow="Agent1 -> Agent2 -> Agent3"
-    )
 
 
-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")
 
 
-def test_add_agent(agent_rearrange):
+# ============================================================================
-    new_agent = MockAgent(name="Agent4")
+# Agent Management Tests
+# ============================================================================
+
+
+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():
+    """Test removing an agent from AgentRearrange."""
+    agents = create_sample_agents()
 
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="ResearchAgent -> WriterAgent -> ReviewerAgent",
+        verbose=True,
+    )
 
-def test_remove_agent(agent_rearrange):
+    agent_rearrange.remove_agent("ReviewerAgent")
-    agent_rearrange.remove_agent("Agent2")
+    assert "ReviewerAgent" not in agent_rearrange.agents
-    assert "Agent2" not in agent_rearrange.agents
+    assert len(agent_rearrange.agents) == 2
+    print("✓ test_remove_agent passed")
 
 
-def test_add_agents(agent_rearrange):
+def test_add_agents():
-    new_agents = [MockAgent(name="Agent4"), MockAgent(name="Agent5")]
+    """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")
+
+
+# ============================================================================
+# 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,
+    )
 
-def test_validate_flow_valid(agent_rearrange):
     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):
+    try:
-    agent_rearrange.flow = "Agent1 -> Agent4"
-    with pytest.raises(ValueError):
         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")
+def test_sequential_flow():
-    assert (
+    """Test sequential flow execution."""
-        result
+    agents = create_sample_agents()
-        == "Agent1 processed Test Task; Agent2 processed Agent1 processed Test Task; Agent3 processed Agent2 processed Agent1 processed Test Task"
+
+    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):
+def test_concurrent_flow():
-    custom_tasks = {"Agent2": "Custom Task"}
+    """Test concurrent flow execution with comma syntax."""
-    result = agent_rearrange.run(
+    agents = create_sample_agents()
-        "Test Task", custom_tasks=custom_tasks
+
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="ResearchAgent, WriterAgent -> ReviewerAgent",
+        max_loops=1,
+        verbose=True,
     )
-    assert (
+
-        result
+    result = agent_rearrange.run("What is 3+3?")
-        == "Agent1 processed Test Task; Agent2 processed Custom Task; Agent3 processed Agent2 processed Custom Task"
+    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):
+def test_get_agent_sequential_awareness():
-    agent_rearrange.human_in_the_loop = True
+    """Test getting agent sequential awareness."""
-    agent_rearrange.custom_human_in_the_loop = MagicMock(
+    agents = create_sample_agents()
-        return_value="Human processed Task"
+
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="ResearchAgent -> WriterAgent -> ReviewerAgent",
+        verbose=True,
     )
-    agent_rearrange.flow = "Agent1 -> H -> Agent3"
+
-    result = agent_rearrange.run("Test Task")
+    awareness = agent_rearrange.get_agent_sequential_awareness("WriterAgent")
-    assert (
+    assert awareness is not None
-        result
+    assert isinstance(awareness, str)
-        == "Agent1 processed Test Task; Human processed Task; Agent3 processed Human processed Task"
+    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):
+def test_run_with_different_output_types():
-    sub_swarm_flow = "Agent1 -> Agent3"
+    """Test run with different output types."""
-    agent_rearrange.add_sub_swarm("SubSwarm1", sub_swarm_flow)
+    agents = create_sample_agents()
-    result = agent_rearrange.run_sub_swarm(
+
-        "SubSwarm1", "Sub Task", None
+    for output_type in ["all", "final"]:
+        agent_rearrange = AgentRearrange(
+            agents=agents,
+            flow="ResearchAgent -> WriterAgent",
+            output_type=output_type,
+            max_loops=1,
+            verbose=True,
         )
-    assert (
+
-        result
+        result = agent_rearrange.run("What is the capital of France?")
-        == "Agent1 processed Sub Task; Agent3 processed Agent1 processed Sub Task"
+        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,
     )
 
+    result = agent_rearrange("What is 10+10?")
+    assert result is not None
+    print("✓ test_callable_execution passed")
 
-def test_process_agent_or_swarm(agent_rearrange):
+
-    result = agent_rearrange.process_agent_or_swarm(
+# ============================================================================
-        "Agent1", "Process Task", None
+# 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,
     )
-    assert result == "Agent1 processed Process Task"
 
+    tasks = ["What is 3+3?", "What is 4+4?"]
+    results = agent_rearrange.concurrent_run(tasks, max_workers=2)
 
-def test_track_history(agent_rearrange):
+    assert results is not None
-    agent_rearrange.track_history("Agent1", "Task Result")
+    assert len(results) == 2
-    assert agent_rearrange.swarm_history["Agent1"] == ["Task Result"]
+    print("✓ test_concurrent_run passed")
 
 
-def test_human_intervention(agent_rearrange):
+# ============================================================================
-    agent_rearrange.human_in_the_loop = True
+# Serialization Tests
-    agent_rearrange.custom_human_in_the_loop = MagicMock(
+# ============================================================================
-        return_value="Human processed Task"
+
+
+def test_to_dict():
+    """Test serialization to dictionary."""
+    agents = create_sample_agents()
+
+    agent_rearrange = AgentRearrange(
+        agents=agents,
+        flow="ResearchAgent -> WriterAgent",
+        verbose=True,
     )
-    result = agent_rearrange.human_intervention("Task")
+
-    assert result == "Human processed 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")
+
+
+# ============================================================================
+# Integration Tests
+# ============================================================================
+
+
+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",
+    )
+
+    # 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"])