|
|
|
|
@ -1,30 +1,41 @@
|
|
|
|
|
import pytest
|
|
|
|
|
from unittest.mock import Mock
|
|
|
|
|
from swarms.structs.utils import find_agent_by_id, find_agent_by_name
|
|
|
|
|
from swarms.structs.agent import Agent
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def create_test_agent(name: str, agent_id: str = None, description: str = "Test agent") -> Agent:
|
|
|
|
|
"""Create a real Agent instance for testing"""
|
|
|
|
|
agent = 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,
|
|
|
|
|
)
|
|
|
|
|
if agent_id:
|
|
|
|
|
agent.id = agent_id
|
|
|
|
|
# Set name attribute for find_agent_by_name
|
|
|
|
|
agent.name = name
|
|
|
|
|
return agent
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_id_found():
|
|
|
|
|
"""Test finding agent by ID when agent exists"""
|
|
|
|
|
mock_agent1 = Mock()
|
|
|
|
|
mock_agent1.id = "agent-1"
|
|
|
|
|
mock_agent1.name = "Agent One"
|
|
|
|
|
|
|
|
|
|
mock_agent2 = Mock()
|
|
|
|
|
mock_agent2.id = "agent-2"
|
|
|
|
|
mock_agent2.name = "Agent Two"
|
|
|
|
|
agent1 = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
agent2 = create_test_agent("Agent Two", "agent-2")
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent1, mock_agent2]
|
|
|
|
|
agents = [agent1, agent2]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_id("agent-1", agents)
|
|
|
|
|
assert result == mock_agent1
|
|
|
|
|
assert result == agent1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_id_not_found():
|
|
|
|
|
"""Test finding agent by ID when agent does not exist"""
|
|
|
|
|
mock_agent1 = Mock()
|
|
|
|
|
mock_agent1.id = "agent-1"
|
|
|
|
|
agent1 = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent1]
|
|
|
|
|
agents = [agent1]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_id("agent-99", agents)
|
|
|
|
|
assert result is None
|
|
|
|
|
@ -32,30 +43,25 @@ def test_find_agent_by_id_not_found():
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_id_with_task():
|
|
|
|
|
"""Test finding agent by ID and running a task"""
|
|
|
|
|
mock_agent = Mock()
|
|
|
|
|
mock_agent.id = "agent-1"
|
|
|
|
|
mock_agent.run.return_value = "Task completed"
|
|
|
|
|
agent = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent]
|
|
|
|
|
agents = [agent]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_id("agent-1", agents, task="Do something")
|
|
|
|
|
assert result == "Task completed"
|
|
|
|
|
mock_agent.run.assert_called_once_with("Do something")
|
|
|
|
|
result = find_agent_by_id("agent-1", agents, task="What is 2+2?")
|
|
|
|
|
assert result is not None
|
|
|
|
|
assert len(str(result)) > 0
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_id_with_task_and_args():
|
|
|
|
|
"""Test finding agent by ID and running a task with args and kwargs"""
|
|
|
|
|
mock_agent = Mock()
|
|
|
|
|
mock_agent.id = "agent-1"
|
|
|
|
|
mock_agent.run.return_value = "Task completed"
|
|
|
|
|
def test_find_agent_by_id_with_task_and_kwargs():
|
|
|
|
|
"""Test finding agent by ID and running a task with kwargs"""
|
|
|
|
|
agent = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent]
|
|
|
|
|
agents = [agent]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_id(
|
|
|
|
|
agent_id="agent-1", agents=agents, task="Do something", kwarg1="value1"
|
|
|
|
|
agent_id="agent-1", agents=agents, task="Say hello"
|
|
|
|
|
)
|
|
|
|
|
assert result == "Task completed"
|
|
|
|
|
mock_agent.run.assert_called_once_with("Do something", kwarg1="value1")
|
|
|
|
|
assert result is not None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_id_empty_list():
|
|
|
|
|
@ -65,39 +71,32 @@ def test_find_agent_by_id_empty_list():
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_id_exception_handling():
|
|
|
|
|
"""Test that find_agent_by_id handles exceptions gracefully"""
|
|
|
|
|
mock_agent = Mock()
|
|
|
|
|
mock_agent.id = "agent-1"
|
|
|
|
|
mock_agent.run.side_effect = Exception("Test error")
|
|
|
|
|
"""Test that find_agent_by_id handles task execution"""
|
|
|
|
|
agent = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent]
|
|
|
|
|
agents = [agent]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_id("agent-1", agents, task="Do something")
|
|
|
|
|
assert result is None
|
|
|
|
|
# Should execute successfully with real agent
|
|
|
|
|
result = find_agent_by_id("agent-1", agents, task="What is the capital of France?")
|
|
|
|
|
assert result is not None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_found():
|
|
|
|
|
"""Test finding agent by name when agent exists"""
|
|
|
|
|
mock_agent1 = Mock()
|
|
|
|
|
mock_agent1.id = "agent-1"
|
|
|
|
|
mock_agent1.name = "Agent One"
|
|
|
|
|
agent1 = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
agent2 = create_test_agent("Agent Two", "agent-2")
|
|
|
|
|
|
|
|
|
|
mock_agent2 = Mock()
|
|
|
|
|
mock_agent2.id = "agent-2"
|
|
|
|
|
mock_agent2.name = "Agent Two"
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent1, mock_agent2]
|
|
|
|
|
agents = [agent1, agent2]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_name("Agent One", agents)
|
|
|
|
|
assert result == mock_agent1
|
|
|
|
|
assert result == agent1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_not_found():
|
|
|
|
|
"""Test finding agent by name when agent does not exist"""
|
|
|
|
|
mock_agent1 = Mock()
|
|
|
|
|
mock_agent1.name = "Agent One"
|
|
|
|
|
agent1 = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent1]
|
|
|
|
|
agents = [agent1]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_name("Agent Ninety Nine", agents)
|
|
|
|
|
assert result is None
|
|
|
|
|
@ -105,30 +104,25 @@ def test_find_agent_by_name_not_found():
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_with_task():
|
|
|
|
|
"""Test finding agent by name and running a task"""
|
|
|
|
|
mock_agent = Mock()
|
|
|
|
|
mock_agent.name = "Agent One"
|
|
|
|
|
mock_agent.run.return_value = "Task completed"
|
|
|
|
|
agent = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent]
|
|
|
|
|
agents = [agent]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_name("Agent One", agents, task="Do something")
|
|
|
|
|
assert result == "Task completed"
|
|
|
|
|
mock_agent.run.assert_called_once_with("Do something")
|
|
|
|
|
result = find_agent_by_name("Agent One", agents, task="What is 3+3?")
|
|
|
|
|
assert result is not None
|
|
|
|
|
assert len(str(result)) > 0
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_with_task_and_args():
|
|
|
|
|
"""Test finding agent by name and running a task with args and kwargs"""
|
|
|
|
|
mock_agent = Mock()
|
|
|
|
|
mock_agent.name = "Agent One"
|
|
|
|
|
mock_agent.run.return_value = "Task completed"
|
|
|
|
|
def test_find_agent_by_name_with_task_and_kwargs():
|
|
|
|
|
"""Test finding agent by name and running a task with kwargs"""
|
|
|
|
|
agent = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent]
|
|
|
|
|
agents = [agent]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_name(
|
|
|
|
|
agent_name="Agent One", agents=agents, task="Do something", kwarg1="value1"
|
|
|
|
|
agent_name="Agent One", agents=agents, task="Say goodbye"
|
|
|
|
|
)
|
|
|
|
|
assert result == "Task completed"
|
|
|
|
|
mock_agent.run.assert_called_once_with("Do something", kwarg1="value1")
|
|
|
|
|
assert result is not None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_empty_list():
|
|
|
|
|
@ -138,23 +132,21 @@ def test_find_agent_by_name_empty_list():
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_name_exception_handling():
|
|
|
|
|
"""Test that find_agent_by_name handles exceptions gracefully"""
|
|
|
|
|
mock_agent = Mock()
|
|
|
|
|
mock_agent.name = "Agent One"
|
|
|
|
|
mock_agent.run.side_effect = Exception("Test error")
|
|
|
|
|
"""Test that find_agent_by_name handles task execution"""
|
|
|
|
|
agent = create_test_agent("Agent One", "agent-1")
|
|
|
|
|
|
|
|
|
|
agents = [mock_agent]
|
|
|
|
|
agents = [agent]
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_name("Agent One", agents, task="Do something")
|
|
|
|
|
assert result is None
|
|
|
|
|
# Should execute successfully with real agent
|
|
|
|
|
result = find_agent_by_name("Agent One", agents, task="List 3 colors")
|
|
|
|
|
assert result is not None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_find_agent_by_id_multiple_agents():
|
|
|
|
|
"""Test finding correct agent by ID when multiple agents exist"""
|
|
|
|
|
agents = []
|
|
|
|
|
for i in range(10):
|
|
|
|
|
agent = Mock()
|
|
|
|
|
agent.id = f"agent-{i}"
|
|
|
|
|
agent = create_test_agent(f"Agent {i}", f"agent-{i}")
|
|
|
|
|
agents.append(agent)
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_id("agent-5", agents)
|
|
|
|
|
@ -165,8 +157,7 @@ def test_find_agent_by_name_multiple_agents():
|
|
|
|
|
"""Test finding correct agent by name when multiple agents exist"""
|
|
|
|
|
agents = []
|
|
|
|
|
for i in range(10):
|
|
|
|
|
agent = Mock()
|
|
|
|
|
agent.name = f"Agent {i}"
|
|
|
|
|
agent = create_test_agent(f"Agent {i}", f"agent-{i}")
|
|
|
|
|
agents.append(agent)
|
|
|
|
|
|
|
|
|
|
result = find_agent_by_name("Agent 5", agents)
|
|
|
|
|
|