|
|
|
@ -1,5 +1,4 @@
|
|
|
|
import pytest
|
|
|
|
import pytest
|
|
|
|
from unittest.mock import Mock, patch, MagicMock
|
|
|
|
|
|
|
|
from swarms.structs.ma_blocks import (
|
|
|
|
from swarms.structs.ma_blocks import (
|
|
|
|
aggregator_agent_task_prompt,
|
|
|
|
aggregator_agent_task_prompt,
|
|
|
|
aggregate,
|
|
|
|
aggregate,
|
|
|
|
@ -7,36 +6,48 @@ from swarms.structs.ma_blocks import (
|
|
|
|
find_agent_by_name,
|
|
|
|
find_agent_by_name,
|
|
|
|
)
|
|
|
|
)
|
|
|
|
from swarms.structs.agent import Agent
|
|
|
|
from swarms.structs.agent import Agent
|
|
|
|
|
|
|
|
from swarms.structs.conversation import Conversation
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def create_test_agent(name: str, description: str = "Test agent") -> Agent:
|
|
|
|
|
|
|
|
"""Create a real Agent instance for testing"""
|
|
|
|
|
|
|
|
return Agent(
|
|
|
|
|
|
|
|
agent_name=name,
|
|
|
|
|
|
|
|
agent_description=description,
|
|
|
|
|
|
|
|
system_prompt=f"You are {name}, a helpful test assistant. Keep responses brief.",
|
|
|
|
|
|
|
|
model_name="gpt-4o-mini",
|
|
|
|
|
|
|
|
max_loops=1,
|
|
|
|
|
|
|
|
verbose=False,
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_aggregator_agent_task_prompt():
|
|
|
|
def test_aggregator_agent_task_prompt():
|
|
|
|
"""Test aggregator agent task prompt generation"""
|
|
|
|
"""Test aggregator agent task prompt generation"""
|
|
|
|
mock_agent1 = Mock()
|
|
|
|
agent1 = create_test_agent("Agent1", "First test agent")
|
|
|
|
mock_agent1.agent_name = "Agent1"
|
|
|
|
agent2 = create_test_agent("Agent2", "Second test agent")
|
|
|
|
|
|
|
|
|
|
|
|
mock_agent2 = Mock()
|
|
|
|
|
|
|
|
mock_agent2.agent_name = "Agent2"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
workers = [mock_agent1, mock_agent2]
|
|
|
|
workers = [agent1, agent2]
|
|
|
|
|
|
|
|
|
|
|
|
mock_conversation = Mock()
|
|
|
|
conversation = Conversation()
|
|
|
|
mock_conversation.get_str.return_value = "Agent1: Hello\nAgent2: Hi"
|
|
|
|
conversation.add(role="Agent1", content="Hello")
|
|
|
|
|
|
|
|
conversation.add(role="Agent2", content="Hi")
|
|
|
|
|
|
|
|
|
|
|
|
result = aggregator_agent_task_prompt(
|
|
|
|
result = aggregator_agent_task_prompt(
|
|
|
|
task="Test task",
|
|
|
|
task="Test task",
|
|
|
|
workers=workers,
|
|
|
|
workers=workers,
|
|
|
|
conversation=mock_conversation
|
|
|
|
conversation=conversation
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
assert "Test task" in result
|
|
|
|
assert "Test task" in result
|
|
|
|
assert "2" in result # Number of agents
|
|
|
|
assert "2" in result # Number of agents
|
|
|
|
assert "Agent1: Hello" in result
|
|
|
|
assert "Hello" in result or "Hi" in result
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_aggregate_missing_task_raises_error():
|
|
|
|
def test_aggregate_missing_task_raises_error():
|
|
|
|
"""Test that missing task raises ValueError"""
|
|
|
|
"""Test that missing task raises ValueError"""
|
|
|
|
|
|
|
|
agent = create_test_agent("TestAgent")
|
|
|
|
with pytest.raises(ValueError, match="Task is required"):
|
|
|
|
with pytest.raises(ValueError, match="Task is required"):
|
|
|
|
aggregate(workers=[Mock()], task=None)
|
|
|
|
aggregate(workers=[agent], task=None)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_aggregate_missing_workers_raises_error():
|
|
|
|
def test_aggregate_missing_workers_raises_error():
|
|
|
|
@ -47,8 +58,9 @@ def test_aggregate_missing_workers_raises_error():
|
|
|
|
|
|
|
|
|
|
|
|
def test_aggregate_workers_not_list_raises_error():
|
|
|
|
def test_aggregate_workers_not_list_raises_error():
|
|
|
|
"""Test that non-list workers raises ValueError"""
|
|
|
|
"""Test that non-list workers raises ValueError"""
|
|
|
|
|
|
|
|
agent = create_test_agent("TestAgent")
|
|
|
|
with pytest.raises(ValueError, match="Workers must be a list"):
|
|
|
|
with pytest.raises(ValueError, match="Workers must be a list"):
|
|
|
|
aggregate(workers=Mock(), task="Test")
|
|
|
|
aggregate(workers=agent, task="Test")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_aggregate_workers_not_callable_raises_error():
|
|
|
|
def test_aggregate_workers_not_callable_raises_error():
|
|
|
|
@ -65,9 +77,9 @@ def test_run_agent_none_agent_raises_error():
|
|
|
|
|
|
|
|
|
|
|
|
def test_run_agent_none_task_raises_error():
|
|
|
|
def test_run_agent_none_task_raises_error():
|
|
|
|
"""Test that None task raises ValueError"""
|
|
|
|
"""Test that None task raises ValueError"""
|
|
|
|
mock_agent = Mock(spec=Agent)
|
|
|
|
agent = create_test_agent("TestAgent")
|
|
|
|
with pytest.raises(ValueError, match="Task cannot be None"):
|
|
|
|
with pytest.raises(ValueError, match="Task cannot be None"):
|
|
|
|
run_agent(agent=mock_agent, task=None)
|
|
|
|
run_agent(agent=agent, task=None)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_run_agent_not_agent_instance_raises_error():
|
|
|
|
def test_run_agent_not_agent_instance_raises_error():
|
|
|
|
@ -78,40 +90,24 @@ def test_run_agent_not_agent_instance_raises_error():
|
|
|
|
|
|
|
|
|
|
|
|
def test_run_agent_success():
|
|
|
|
def test_run_agent_success():
|
|
|
|
"""Test successful agent run"""
|
|
|
|
"""Test successful agent run"""
|
|
|
|
mock_agent = Mock(spec=Agent)
|
|
|
|
agent = create_test_agent("TestAgent")
|
|
|
|
mock_agent.run.return_value = "Task completed"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
result = run_agent(agent=mock_agent, task="Test task")
|
|
|
|
result = run_agent(agent=agent, task="What is 2+2?")
|
|
|
|
|
|
|
|
|
|
|
|
assert result == "Task completed"
|
|
|
|
assert result is not None
|
|
|
|
mock_agent.run.assert_called_once_with(task="Test task")
|
|
|
|
assert len(str(result)) > 0
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_run_agent_with_args_kwargs():
|
|
|
|
def test_run_agent_with_kwargs():
|
|
|
|
"""Test run_agent with additional args and kwargs"""
|
|
|
|
"""Test run_agent with additional kwargs"""
|
|
|
|
mock_agent = Mock(spec=Agent)
|
|
|
|
agent = create_test_agent("TestAgent")
|
|
|
|
mock_agent.run.return_value = "Success"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
result = run_agent(
|
|
|
|
result = run_agent(
|
|
|
|
agent=mock_agent,
|
|
|
|
agent=agent,
|
|
|
|
task="Test",
|
|
|
|
task="Say hello"
|
|
|
|
extra_param="value"
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
assert result == "Success"
|
|
|
|
|
|
|
|
mock_agent.run.assert_called_once_with(
|
|
|
|
|
|
|
|
task="Test",
|
|
|
|
|
|
|
|
extra_param="value"
|
|
|
|
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
assert result is not None
|
|
|
|
def test_run_agent_runtime_error_on_exception():
|
|
|
|
|
|
|
|
"""Test that exceptions during run raise RuntimeError"""
|
|
|
|
|
|
|
|
mock_agent = Mock(spec=Agent)
|
|
|
|
|
|
|
|
mock_agent.run.side_effect = Exception("Agent failed")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
with pytest.raises(RuntimeError, match="Error running agent"):
|
|
|
|
|
|
|
|
run_agent(agent=mock_agent, task="Test")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_empty_list_raises_error():
|
|
|
|
def test_find_agent_by_name_empty_list_raises_error():
|
|
|
|
@ -122,52 +118,56 @@ def test_find_agent_by_name_empty_list_raises_error():
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_non_string_raises_error():
|
|
|
|
def test_find_agent_by_name_non_string_raises_error():
|
|
|
|
"""Test that non-string agent_name raises TypeError"""
|
|
|
|
"""Test that non-string agent_name raises TypeError"""
|
|
|
|
mock_agent = Mock()
|
|
|
|
agent = create_test_agent("TestAgent")
|
|
|
|
with pytest.raises(TypeError, match="Agent name must be a string"):
|
|
|
|
with pytest.raises(TypeError, match="Agent name must be a string"):
|
|
|
|
find_agent_by_name(agents=[mock_agent], agent_name=123)
|
|
|
|
find_agent_by_name(agents=[agent], agent_name=123)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_empty_string_raises_error():
|
|
|
|
def test_find_agent_by_name_empty_string_raises_error():
|
|
|
|
"""Test that empty agent_name raises ValueError"""
|
|
|
|
"""Test that empty agent_name raises ValueError"""
|
|
|
|
mock_agent = Mock()
|
|
|
|
agent = create_test_agent("TestAgent")
|
|
|
|
with pytest.raises(ValueError, match="Agent name cannot be empty"):
|
|
|
|
with pytest.raises(ValueError, match="Agent name cannot be empty"):
|
|
|
|
find_agent_by_name(agents=[mock_agent], agent_name=" ")
|
|
|
|
find_agent_by_name(agents=[agent], agent_name=" ")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_success():
|
|
|
|
def test_find_agent_by_name_success():
|
|
|
|
"""Test successful agent finding by name"""
|
|
|
|
"""Test successful agent finding by name"""
|
|
|
|
mock_agent1 = Mock()
|
|
|
|
agent1 = create_test_agent("Agent1")
|
|
|
|
mock_agent1.name = "Agent1"
|
|
|
|
agent2 = create_test_agent("Agent2")
|
|
|
|
|
|
|
|
|
|
|
|
mock_agent2 = Mock()
|
|
|
|
# Note: find_agent_by_name looks for 'name' attribute, not 'agent_name'
|
|
|
|
mock_agent2.name = "Agent2"
|
|
|
|
agent1.name = "Agent1"
|
|
|
|
|
|
|
|
agent2.name = "Agent2"
|
|
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_name(
|
|
|
|
result = find_agent_by_name(
|
|
|
|
agents=[mock_agent1, mock_agent2],
|
|
|
|
agents=[agent1, agent2],
|
|
|
|
agent_name="Agent2"
|
|
|
|
agent_name="Agent2"
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
assert result == mock_agent2
|
|
|
|
assert result == agent2
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_not_found_raises_error():
|
|
|
|
def test_find_agent_by_name_not_found_raises_error():
|
|
|
|
"""Test that agent not found raises RuntimeError"""
|
|
|
|
"""Test that agent not found raises RuntimeError"""
|
|
|
|
mock_agent = Mock()
|
|
|
|
agent = create_test_agent("Agent1")
|
|
|
|
mock_agent.name = "Agent1"
|
|
|
|
agent.name = "Agent1"
|
|
|
|
|
|
|
|
|
|
|
|
with pytest.raises(RuntimeError, match="Error finding agent"):
|
|
|
|
with pytest.raises(RuntimeError, match="Error finding agent"):
|
|
|
|
find_agent_by_name(agents=[mock_agent], agent_name="NonExistent")
|
|
|
|
find_agent_by_name(agents=[agent], agent_name="NonExistent")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_agent_with_name_attribute():
|
|
|
|
|
|
|
|
"""Test finding agent when agent has name attribute"""
|
|
|
|
|
|
|
|
agent1 = create_test_agent("Agent1")
|
|
|
|
|
|
|
|
agent2 = create_test_agent("TargetAgent")
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_agent_without_name_attribute():
|
|
|
|
# Set the name attribute that find_agent_by_name looks for
|
|
|
|
"""Test finding agent when some agents don't have name attribute"""
|
|
|
|
agent1.name = "Agent1"
|
|
|
|
mock_agent1 = Mock(spec=[]) # No name attribute
|
|
|
|
agent2.name = "TargetAgent"
|
|
|
|
mock_agent2 = Mock()
|
|
|
|
|
|
|
|
mock_agent2.name = "TargetAgent"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_name(
|
|
|
|
result = find_agent_by_name(
|
|
|
|
agents=[mock_agent1, mock_agent2],
|
|
|
|
agents=[agent1, agent2],
|
|
|
|
agent_name="TargetAgent"
|
|
|
|
agent_name="TargetAgent"
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
assert result == mock_agent2
|
|
|
|
assert result == agent2
|
|
|
|
|