import os
import traceback
from datetime import datetime
from typing import Callable, Dict, List, Optional
from loguru import logger
from swarm_models import OpenAIChat
from swarms.structs.agent import Agent
from swarms.structs.rearrange import AgentRearrange
class TestResult:
"""Class to store test results and metadata"""
def __init__(self, test_name: str):
self.test_name = test_name
self.start_time = datetime.now()
self.end_time = None
self.success = False
self.error = None
self.traceback = None
self.function_output = None
def complete(
self, success: bool, error: Optional[Exception] = None
):
"""Complete the test execution with results"""
self.end_time = datetime.now()
self.success = success
if error:
self.error = str(error)
self.traceback = traceback.format_exc()
def duration(self) -> float:
"""Calculate test duration in seconds"""
if self.end_time:
return (self.end_time - self.start_time).total_seconds()
return 0
def run_test(test_func: Callable) -> TestResult:
"""
Decorator to run tests with error handling and logging
Args:
test_func (Callable): Test function to execute
Returns:
TestResult: Object containing test execution details
"""
def wrapper(*args, **kwargs) -> TestResult:
result = TestResult(test_func.__name__)
logger.info(
f"\n{'='*20} Running test: {test_func.__name__} {'='*20}"
)
try:
output = test_func(*args, **kwargs)
result.function_output = output
result.complete(success=True)
logger.success(
f"ā
Test {test_func.__name__} passed successfully"
)
except Exception as e:
result.complete(success=False, error=e)
logger.error(
f"ā Test {test_func.__name__} failed with error: {str(e)}"
)
logger.error(f"Traceback: {traceback.format_exc()}")
logger.info(
f"Test duration: {result.duration():.2f} seconds\n"
)
return result
return wrapper
def create_functional_agents() -> List[Agent]:
"""
Create a list of functional agents with real LLM integration for testing.
Using OpenAI's GPT model for realistic agent behavior testing.
"""
# Initialize OpenAI Chat model
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
logger.warning(
"No OpenAI API key found. Using mock agents instead."
)
return [
create_mock_agent("TestAgent1"),
create_mock_agent("TestAgent2"),
]
try:
model = OpenAIChat(
api_key=api_key, model_name="gpt-4o", temperature=0.1
)
# Create boss agent
boss_agent = Agent(
agent_name="BossAgent",
system_prompt="""
You are the BossAgent responsible for managing and overseeing test scenarios.
Your role is to coordinate tasks between agents and ensure efficient collaboration.
Analyze inputs, break down tasks, and provide clear directives to other agents.
Maintain a structured approach to task management and result compilation.
""",
llm=model,
max_loops=1,
dashboard=False,
streaming_on=True,
verbose=True,
stopping_token="<DONE>",
state_save_file_type="json",
saved_state_path="test_boss_agent.json",
)
# Create analysis agent
analysis_agent = Agent(
agent_name="AnalysisAgent",
system_prompt="""
You are the AnalysisAgent responsible for detailed data processing and analysis.
Your role is to examine input data, identify patterns, and provide analytical insights.
Focus on breaking down complex information into clear, actionable components.
""",
llm=model,
max_loops=1,
dashboard=False,
streaming_on=True,
verbose=True,
stopping_token="<DONE>",
state_save_file_type="json",
saved_state_path="test_analysis_agent.json",
)
# Create summary agent
summary_agent = Agent(
agent_name="SummaryAgent",
system_prompt="""
You are the SummaryAgent responsible for consolidating and summarizing information.
Your role is to take detailed analysis and create concise, actionable summaries.
Focus on highlighting key points and ensuring clarity in communication.
""",
llm=model,
max_loops=1,
dashboard=False,
streaming_on=True,
verbose=True,
stopping_token="<DONE>",
state_save_file_type="json",
saved_state_path="test_summary_agent.json",
)
logger.info(
"Successfully created functional agents with LLM integration"
)
return [boss_agent, analysis_agent, summary_agent]
except Exception as e:
logger.error(f"Failed to create functional agents: {str(e)}")
logger.warning("Falling back to mock agents")
return [
create_mock_agent("TestAgent1"),
create_mock_agent("TestAgent2"),
]
def create_mock_agent(name: str) -> Agent:
"""Create a mock agent for testing when LLM integration is not available"""
return Agent(
agent_name=name,
system_prompt=f"You are a test agent named {name}",
llm=None,
)
@run_test
def test_init():
"""Test AgentRearrange initialization with functional agents"""
logger.info("Creating agents for initialization test")
agents = create_functional_agents()
rearrange = AgentRearrange(
name="TestRearrange",
agents=agents,
flow=f"{agents[0].agent_name} -> {agents[1].agent_name} -> {agents[2].agent_name}",
)
assert rearrange.name == "TestRearrange"
assert len(rearrange.agents) == 3
assert (
rearrange.flow
== f"{agents[0].agent_name} -> {agents[1].agent_name} -> {agents[2].agent_name}"
)
logger.info(
f"Initialized AgentRearrange with {len(agents)} agents"
)
return True
@run_test
def test_validate_flow():
"""Test flow validation logic"""
agents = create_functional_agents()
rearrange = AgentRearrange(
agents=agents,
flow=f"{agents[0].agent_name} -> {agents[1].agent_name}",
)
logger.info("Testing valid flow pattern")
valid = rearrange.validate_flow()
assert valid is True
logger.info("Testing invalid flow pattern")
rearrange.flow = f"{agents[0].agent_name} {agents[1].agent_name}" # Missing arrow
try:
rearrange.validate_flow()
assert False, "Should have raised ValueError"
except ValueError as e:
logger.info(
f"Successfully caught invalid flow error: {str(e)}"
)
assert True
return True
@run_test
def test_add_remove_agent():
"""Test adding and removing agents from the swarm"""
agents = create_functional_agents()
rearrange = AgentRearrange(
agents=agents[:2]
) # Start with first two agents
logger.info("Testing agent addition")
new_agent = agents[2] # Use the third agent as new agent
rearrange.add_agent(new_agent)
assert new_agent.agent_name in rearrange.agents
logger.info("Testing agent removal")
rearrange.remove_agent(new_agent.agent_name)
assert new_agent.agent_name not in rearrange.agents
return True
@run_test
def test_basic_run():
"""Test basic task execution with the swarm"""
agents = create_functional_agents()
rearrange = AgentRearrange(
name="TestSwarm",
agents=agents,
flow=f"{agents[0].agent_name} -> {agents[1].agent_name} -> {agents[2].agent_name}",
max_loops=1,
)
test_task = (
"Analyze this test message and provide a brief summary."
)
logger.info(f"Running test task: {test_task}")
try:
result = rearrange.run(test_task)
assert result is not None
logger.info(
f"Successfully executed task with result length: {len(str(result))}"
)
return True
except Exception as e:
logger.error(f"Task execution failed: {str(e)}")
raise
def run_all_tests() -> Dict[str, TestResult]:
"""
Run all test cases and collect results
Returns:
Dict[str, TestResult]: Dictionary mapping test names to their results
"""
logger.info("\nš Starting AgentRearrange test suite execution")
test_functions = [
test_init,
test_validate_flow,
test_add_remove_agent,
test_basic_run,
]
results = {}
for test in test_functions:
result = test()
results[test.__name__] = result
# Log summary
total_tests = len(results)
passed_tests = sum(1 for r in results.values() if r.success)
failed_tests = total_tests - passed_tests
logger.info("\nš Test Suite Summary:")
logger.info(f"Total Tests: {total_tests}")
print(f"ā
Passed: {passed_tests}")
if failed_tests > 0:
logger.error(f"ā Failed: {failed_tests}")
# Detailed failure information
if failed_tests > 0:
logger.error("\nā Failed Tests Details:")
for name, result in results.items():
if not result.success:
logger.error(f"\n{name}:")
logger.error(f"Error: {result.error}")
logger.error(f"Traceback: {result.traceback}")
return results
if __name__ == "__main__":
print("š Starting AgentRearrange Test Suite")
results = run_all_tests()
print("š Test Suite Execution Completed")