swarms/tests/structs/test_utils.py

import pytest
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"""
    agent1 = create_test_agent("Agent One", "agent-1")
    agent2 = create_test_agent("Agent Two", "agent-2")

    agents = [agent1, agent2]

    result = find_agent_by_id("agent-1", agents)
    assert result == agent1


def test_find_agent_by_id_not_found():
    """Test finding agent by ID when agent does not exist"""
    agent1 = create_test_agent("Agent One", "agent-1")

    agents = [agent1]

    result = find_agent_by_id("agent-99", agents)
    assert result is None


def test_find_agent_by_id_with_task():
    """Test finding agent by ID and running a task"""
    agent = create_test_agent("Agent One", "agent-1")

    agents = [agent]

    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_kwargs():
    """Test finding agent by ID and running a task with kwargs"""
    agent = create_test_agent("Agent One", "agent-1")

    agents = [agent]

    result = find_agent_by_id(
        agent_id="agent-1", agents=agents, task="Say hello"
    )
    assert result is not None


def test_find_agent_by_id_empty_list():
    """Test finding agent by ID in empty list"""
    result = find_agent_by_id("agent-1", [])
    assert result is None


def test_find_agent_by_id_exception_handling():
    """Test that find_agent_by_id handles task execution"""
    agent = create_test_agent("Agent One", "agent-1")

    agents = [agent]

    # 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"""
    agent1 = create_test_agent("Agent One", "agent-1")
    agent2 = create_test_agent("Agent Two", "agent-2")

    agents = [agent1, agent2]

    result = find_agent_by_name("Agent One", agents)
    assert result == agent1


def test_find_agent_by_name_not_found():
    """Test finding agent by name when agent does not exist"""
    agent1 = create_test_agent("Agent One", "agent-1")

    agents = [agent1]

    result = find_agent_by_name("Agent Ninety Nine", agents)
    assert result is None


def test_find_agent_by_name_with_task():
    """Test finding agent by name and running a task"""
    agent = create_test_agent("Agent One", "agent-1")

    agents = [agent]

    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_kwargs():
    """Test finding agent by name and running a task with kwargs"""
    agent = create_test_agent("Agent One", "agent-1")

    agents = [agent]

    result = find_agent_by_name(
        agent_name="Agent One", agents=agents, task="Say goodbye"
    )
    assert result is not None


def test_find_agent_by_name_empty_list():
    """Test finding agent by name in empty list"""
    result = find_agent_by_name("Agent One", [])
    assert result is None


def test_find_agent_by_name_exception_handling():
    """Test that find_agent_by_name handles task execution"""
    agent = create_test_agent("Agent One", "agent-1")

    agents = [agent]

    # 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 = create_test_agent(f"Agent {i}", f"agent-{i}")
        agents.append(agent)

    result = find_agent_by_id("agent-5", agents)
    assert result.id == "agent-5"


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 = create_test_agent(f"Agent {i}", f"agent-{i}")
        agents.append(agent)

    result = find_agent_by_name("Agent 5", agents)
    assert result.name == "Agent 5"