Merge pull request #1172 from Steve-Dusty/master
align tests with actual SwarmRouter API for test_swarm_routerpull/1173/head
Kye Gomez
1 week ago
committed by
GitHub
commit
efe293df9e
@ -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")
|
||||
|
||||
|
||||
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("Agent2")
|
||||
assert "Agent2" not in agent_rearrange.agents
|
||||
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")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 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):
|
||||
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,
|
||||
)
|
||||
assert (
|
||||
result
|
||||
== "Agent1 processed Sub Task; Agent3 processed Agent1 processed Sub Task"
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
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):
|
||||
agent_rearrange.track_history("Agent1", "Task Result")
|
||||
assert agent_rearrange.swarm_history["Agent1"] == ["Task Result"]
|
||||
assert results is not None
|
||||
assert len(results) == 2
|
||||
print("✓ test_concurrent_run passed")
|
||||
|
||||
|
||||
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"
|
||||
# ============================================================================
|
||||
# Serialization Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
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"])
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Loading…
Reference in new issue