Skip to content

Commit

Permalink
fix test
Browse files Browse the repository at this point in the history
  • Loading branch information
bhancockio committed Oct 23, 2024
1 parent 468ebae commit 5cfafa7
Showing 1 changed file with 10 additions and 8 deletions.
18 changes: 10 additions & 8 deletions tests/crew_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -1727,10 +1727,13 @@ def testing_tool(first_number: int, second_number: int) -> int:
crew.kickoff()


@patch("crewai.crew.Crew.kickoff")
@patch("crewai.crew.CrewTrainingHandler")
@patch("crewai.crew.TaskEvaluator")
@patch("crewai.crew.Crew.copy")
def test_crew_train_success(copy_mock, task_evaluator, crew_training_handler):
def test_crew_train_success(
copy_mock, task_evaluator, crew_training_handler, kickoff_mock
):
task = Task(
description="Come up with a list of 5 interesting ideas to explore for an article, then write one amazing paragraph highlight for each idea that showcases how good an article about this topic could be. Return the list of ideas with their paragraph and your notes.",
expected_output="5 bullet points with a paragraph for each idea.",
Expand All @@ -1743,15 +1746,14 @@ def test_crew_train_success(copy_mock, task_evaluator, crew_training_handler):
)

# Create a mock for the copied crew
copied_crew = mock.Mock()
copy_mock.return_value = copied_crew
copy_mock.return_value = crew

crew.train(
n_iterations=2, inputs={"topic": "AI"}, filename="trained_agents_data.pkl"
)

# Ensure kickoff is called on the copied crew
copied_crew.kickoff.assert_has_calls(
kickoff_mock.assert_has_calls(
[mock.call(inputs={"topic": "AI"}), mock.call(inputs={"topic": "AI"})]
)

Expand Down Expand Up @@ -2504,7 +2506,8 @@ def test_conditional_should_execute():

@mock.patch("crewai.crew.CrewEvaluator")
@mock.patch("crewai.crew.Crew.copy")
def test_crew_testing_function(copy_mock, crew_evaluator):
@mock.patch("crewai.crew.Crew.kickoff")
def test_crew_testing_function(kickoff_mock, copy_mock, crew_evaluator):
task = Task(
description="Come up with a list of 5 interesting ideas to explore for an article, then write one amazing paragraph highlight for each idea that showcases how good an article about this topic could be. Return the list of ideas with their paragraph and your notes.",
expected_output="5 bullet points with a paragraph for each idea.",
Expand All @@ -2517,14 +2520,13 @@ def test_crew_testing_function(copy_mock, crew_evaluator):
)

# Create a mock for the copied crew
copied_crew = mock.Mock()
copy_mock.return_value = copied_crew
copy_mock.return_value = crew

n_iterations = 2
crew.test(n_iterations, openai_model_name="gpt-4o-mini", inputs={"topic": "AI"})

# Ensure kickoff is called on the copied crew
copied_crew.kickoff.assert_has_calls(
kickoff_mock.assert_has_calls(
[mock.call(inputs={"topic": "AI"}), mock.call(inputs={"topic": "AI"})]
)

Expand Down

0 comments on commit 5cfafa7

Please sign in to comment.