From 540e328f0647749be44159e1bddd54888dd46581 Mon Sep 17 00:00:00 2001
From: Eduardo Chiarotti <dudumelgaco@hotmail.com>
Date: Thu, 8 Aug 2024 21:06:10 -0300
Subject: [PATCH] feat: fix tests

---
 tests/crew_test.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/crew_test.py b/tests/crew_test.py
index b9e9a1ca6..62fd8b7a3 100644
--- a/tests/crew_test.py
+++ b/tests/crew_test.py
@@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch
 
 import pydantic_core
 import pytest
+
 from crewai.agent import Agent
 from crewai.agents.cache import CacheHandler
 from crewai.crew import Crew
@@ -1798,7 +1799,9 @@ def test_crew_train_success(task_evaluator, crew_training_handler, kickoff):
         agents=[researcher, writer],
         tasks=[task],
     )
-    crew.train(n_iterations=2, inputs={"topic": "AI"})
+    crew.train(
+        n_iterations=2, inputs={"topic": "AI"}, filename="trained_agents_data.pkl"
+    )
     task_evaluator.assert_has_calls(
         [
             mock.call(researcher),
@@ -1882,7 +1885,7 @@ def test__setup_for_training():
     for agent in agents:
         assert agent.allow_delegation is True
 
-    crew._setup_for_training()
+    crew._setup_for_training("trained_agents_data.pkl")
 
     assert crew._train is True
     assert task.human_input is True