mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 16:48:30 +00:00
96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
import pytest
|
|
|
|
from crewai import Agent, Crew, Task
|
|
from crewai.tasks.conditional_task import ConditionalTask
|
|
from crewai.tasks.task_output import TaskOutput
|
|
|
|
|
|
def test_conditional_task_preserved_in_copy():
|
|
"""Test that ConditionalTask objects are preserved when copying a Crew."""
|
|
agent = Agent(
|
|
role="Researcher",
|
|
goal="Research topics",
|
|
backstory="You are a researcher."
|
|
)
|
|
|
|
# Create a regular task
|
|
task1 = Task(
|
|
description="Research topic A",
|
|
expected_output="Research results for topic A",
|
|
agent=agent
|
|
)
|
|
|
|
# Create a conditional task
|
|
conditional_task = ConditionalTask(
|
|
description="Research topic B if topic A was successful",
|
|
expected_output="Research results for topic B",
|
|
agent=agent,
|
|
condition=lambda output: "success" in output.raw.lower()
|
|
)
|
|
|
|
# Create a crew with both tasks
|
|
crew = Crew(
|
|
agents=[agent],
|
|
tasks=[task1, conditional_task]
|
|
)
|
|
|
|
# Create a copy of the crew
|
|
crew_copy = crew.copy()
|
|
|
|
# Check that the conditional task is still a ConditionalTask in the copied crew
|
|
assert isinstance(crew_copy.tasks[1], ConditionalTask)
|
|
assert hasattr(crew_copy.tasks[1], "should_execute")
|
|
|
|
def test_conditional_task_preserved_in_kickoff_for_each():
|
|
"""Test that ConditionalTask objects are preserved when using kickoff_for_each."""
|
|
from unittest.mock import patch
|
|
|
|
agent = Agent(
|
|
role="Researcher",
|
|
goal="Research topics",
|
|
backstory="You are a researcher."
|
|
)
|
|
|
|
# Create a regular task
|
|
task1 = Task(
|
|
description="Research topic A",
|
|
expected_output="Research results for topic A",
|
|
agent=agent
|
|
)
|
|
|
|
# Create a conditional task
|
|
conditional_task = ConditionalTask(
|
|
description="Research topic B if topic A was successful",
|
|
expected_output="Research results for topic B",
|
|
agent=agent,
|
|
condition=lambda output: "success" in output.raw.lower()
|
|
)
|
|
|
|
# Create a crew with both tasks
|
|
crew = Crew(
|
|
agents=[agent],
|
|
tasks=[task1, conditional_task]
|
|
)
|
|
|
|
# Mock the kickoff method to avoid actual execution
|
|
with patch.object(Crew, "kickoff") as mock_kickoff:
|
|
# Set up the mock to return a TaskOutput
|
|
mock_output = TaskOutput(
|
|
description="Mock task output",
|
|
raw="Success with topic",
|
|
agent=agent.role
|
|
)
|
|
mock_kickoff.return_value = mock_output
|
|
|
|
# Call kickoff_for_each with test inputs
|
|
inputs = [{"topic": "test1"}, {"topic": "test2"}]
|
|
crew.kickoff_for_each(inputs=inputs)
|
|
|
|
# Verify the mock was called with the expected inputs
|
|
assert mock_kickoff.call_count == len(inputs)
|
|
|
|
# Create a copy of the crew to verify the type preservation
|
|
# (since we can't directly access the crews created inside kickoff_for_each)
|
|
crew_copy = crew.copy()
|
|
assert isinstance(crew_copy.tasks[1], ConditionalTask)
|