From 468ebae5e69b54d59319663d3bf6a6ca06a38880 Mon Sep 17 00:00:00 2001 From: Brandon Hancock Date: Tue, 22 Oct 2024 21:13:36 -0400 Subject: [PATCH] fix last test --- tests/crew_test.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/tests/crew_test.py b/tests/crew_test.py index 78cabeeaa..9e2f05219 100644 --- a/tests/crew_test.py +++ b/tests/crew_test.py @@ -1727,11 +1727,10 @@ def test_manager_agent_with_tools_raises_exception(): 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, kickoff): +def test_crew_train_success(copy_mock, task_evaluator, crew_training_handler): 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.", @@ -1744,14 +1743,15 @@ def test_crew_train_success(copy_mock, task_evaluator, crew_training_handler, ki ) # Create a mock for the copied crew - copy_mock.return_value = crew + copied_crew = mock.Mock() + copy_mock.return_value = copied_crew crew.train( n_iterations=2, inputs={"topic": "AI"}, filename="trained_agents_data.pkl" ) # Ensure kickoff is called on the copied crew - crew.kickoff.assert_has_calls( + copied_crew.kickoff.assert_has_calls( [mock.call(inputs={"topic": "AI"}), mock.call(inputs={"topic": "AI"})] ) @@ -2503,8 +2503,8 @@ def test_conditional_should_execute(): @mock.patch("crewai.crew.CrewEvaluator") -@mock.patch("crewai.crew.Crew.kickoff") -def test_crew_testing_function(mock_kickoff, crew_evaluator): +@mock.patch("crewai.crew.Crew.copy") +def test_crew_testing_function(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.", @@ -2515,11 +2515,16 @@ def test_crew_testing_function(mock_kickoff, crew_evaluator): agents=[researcher], tasks=[task], ) + + # Create a mock for the copied crew + copied_crew = mock.Mock() + copy_mock.return_value = copied_crew + n_iterations = 2 crew.test(n_iterations, openai_model_name="gpt-4o-mini", inputs={"topic": "AI"}) - assert len(mock_kickoff.mock_calls) == n_iterations - mock_kickoff.assert_has_calls( + # Ensure kickoff is called on the copied crew + copied_crew.kickoff.assert_has_calls( [mock.call(inputs={"topic": "AI"}), mock.call(inputs={"topic": "AI"})] )