crewAI/tests/procedure/test_procedure.py

from unittest.mock import MagicMock

import pytest

from crewai.agent import Agent
from crewai.crew import Crew
from crewai.crews.crew_output import CrewOutput
from crewai.procedure.procedure import Procedure
from crewai.process import Process
from crewai.task import Task
from crewai.tasks.task_output import TaskOutput


@pytest.fixture
def mock_crew_factory():
    def _create_mock_crew():
        crew = MagicMock(spec=Crew)
        task_output = TaskOutput(
            description="Test task", raw="Task output", agent="Test Agent"
        )
        crew_output = CrewOutput(
            raw="Test output",
            tasks_output=[task_output],
            token_usage={
                "total_tokens": 100,
                "prompt_tokens": 50,
                "completion_tokens": 50,
            },
            json_dict={"key": "value"},
        )

        async def async_kickoff(inputs=None):
            return crew_output

        crew.kickoff.return_value = crew_output
        crew.kickoff_async.side_effect = async_kickoff

        # Add more attributes that Procedure might be expecting
        crew.verbose = 0
        crew.output_log_file = None
        crew.max_rpm = None
        crew.memory = False
        crew.process = Process.sequential
        crew.config = None
        crew.cache = True

        # Add non-empty agents and tasks
        mock_agent = MagicMock(spec=Agent)
        mock_task = MagicMock(spec=Task)
        mock_task.agent = mock_agent
        mock_task.async_execution = False
        mock_task.context = None

        crew.agents = [mock_agent]
        crew.tasks = [mock_task]

        return crew

    return _create_mock_crew


def test_procedure_initialization(mock_crew_factory):
    """
    Test that a Procedure is correctly initialized with the given crews.
    """
    crew1 = mock_crew_factory()
    crew2 = mock_crew_factory()

    procedure = Procedure(crews=[crew1, crew2])
    assert len(procedure.crews) == 2
    assert procedure.crews[0] == crew1
    assert procedure.crews[1] == crew2


@pytest.mark.asyncio
async def test_procedure_kickoff_single_input(mock_crew_factory):
    """
    Test that Procedure.kickoff() correctly processes a single input
    and returns the expected CrewOutput.
    """
    mock_crew_1 = mock_crew_factory()
    procedure = Procedure(crews=[mock_crew_1])
    input_data = {"key": "value"}
    result = await procedure.kickoff([input_data])

    mock_crew_1.kickoff_async.assert_called_once_with(inputs=input_data)
    assert len(result) == 1
    assert isinstance(result[0], CrewOutput)
    assert result[0].raw == "Test output"
    assert len(result[0].tasks_output) == 1
    assert result[0].tasks_output[0].raw == "Task output"
    assert result[0].token_usage == {
        "total_tokens": 100,
        "prompt_tokens": 50,
        "completion_tokens": 50,
    }


@pytest.mark.asyncio
async def test_procedure_kickoff_multiple_inputs(mock_crew_factory):
    """
    Test that Procedure.kickoff() correctly processes multiple inputs
    and returns the expected CrewOutputs.
    """
    mock_crew_1, mock_crew_2 = mock_crew_factory(), mock_crew_factory()
    procedure = Procedure(crews=[mock_crew_1, mock_crew_2])
    input_data = [{"key1": "value1"}, {"key2": "value2"}]
    result = await procedure.kickoff(input_data)

    expected_call_count_per_crew = 2
    assert mock_crew_1.kickoff_async.call_count == expected_call_count_per_crew
    assert mock_crew_2.kickoff_async.call_count == expected_call_count_per_crew
    assert len(result) == 2
    assert all(isinstance(r, CrewOutput) for r in result)
    assert all(len(r.tasks_output) == 1 for r in result)
    assert all(
        r.token_usage
        == {"total_tokens": 100, "prompt_tokens": 50, "completion_tokens": 50}
        for r in result
    )


@pytest.mark.asyncio
async def test_procedure_chaining(mock_crew_factory):
    """
    Test that Procedure correctly chains multiple crews, passing the output
    of one crew as input to the next crew in the sequence.

    This test verifies:
    1. The first crew receives the initial input.
    2. The second crew receives the output from the first crew as its input.
    3. The final output contains the result from the last crew in the chain.
    4. Task outputs and token usage are correctly propagated through the chain.
    """
    crew1, crew2 = mock_crew_factory(), mock_crew_factory()
    task_output1 = TaskOutput(description="Task 1", raw="Output 1", agent="Agent 1")
    task_output2 = TaskOutput(description="Task 2", raw="Final output", agent="Agent 2")

    crew_output1 = CrewOutput(
        raw="Output 1",
        tasks_output=[task_output1],
        token_usage={"total_tokens": 100, "prompt_tokens": 50, "completion_tokens": 50},
        json_dict={"key1": "value1"},
    )
    crew_output2 = CrewOutput(
        raw="Final output",
        tasks_output=[task_output2],
        token_usage={"total_tokens": 150, "prompt_tokens": 75, "completion_tokens": 75},
        json_dict={"key2": "value2"},
    )

    async def async_kickoff1(inputs=None):
        return crew_output1

    async def async_kickoff2(inputs=None):
        return crew_output2

    crew1.kickoff_async.side_effect = async_kickoff1
    crew2.kickoff_async.side_effect = async_kickoff2

    procedure = Procedure(crews=[crew1, crew2])
    input_data = [{"initial": "data"}]
    result = await procedure.kickoff(input_data)

    # Check that the first crew received the initial input
    crew1.kickoff_async.assert_called_once_with(inputs={"initial": "data"})

    # Check that the second crew received the output from the first crew as its input
    crew2.kickoff_async.assert_called_once_with(inputs=crew_output1.to_dict())

    # Check the final output
    assert len(result) == 1
    assert isinstance(result[0], CrewOutput)
    assert result[0].raw == "Final output"
    assert len(result[0].tasks_output) == 1
    assert result[0].tasks_output[0].raw == "Final output"
    assert result[0].token_usage == {
        "total_tokens": 150,
        "prompt_tokens": 75,
        "completion_tokens": 75,
    }
    assert result[0].json_dict == {"key2": "value2"}


def test_crew_rshift_operator():
    """
    Test that the >> operator correctly creates a Procedure from two Crews.
    """
    # Create minimal Crew instances
    agent = Agent(role="Test Agent", goal="Test Goal", backstory="Test Backstory")
    task = Task(agent=agent, description="Test Task", expected_output="Test Output")
    crew1 = Crew(agents=[agent], tasks=[task])
    crew2 = Crew(agents=[agent], tasks=[task])
    crew3 = Crew(agents=[agent], tasks=[task])

    # Test the >> operator
    procedure = crew1 >> crew2

    assert isinstance(procedure, Procedure)
    assert len(procedure.crews) == 2
    assert procedure.crews[0] == crew1
    assert procedure.crews[1] == crew2

    # Test chaining multiple crews
    procedure = crew1 >> crew2 >> crew3

    assert isinstance(procedure, Procedure)
    assert len(procedure.crews) == 3
    assert procedure.crews[0] == crew1
    assert procedure.crews[1] == crew2
    assert procedure.crews[2] == crew3

    # Test error case: trying to shift with non-Crew object
    with pytest.raises(TypeError):
        crew1 >> "not a crew"