feat: fix tests

This commit is contained in:
Eduardo Chiarotti
2024-08-08 21:06:10 -03:00
parent 87d4c5f092
commit 540e328f06

View File

@@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch
import pydantic_core import pydantic_core
import pytest import pytest
from crewai.agent import Agent from crewai.agent import Agent
from crewai.agents.cache import CacheHandler from crewai.agents.cache import CacheHandler
from crewai.crew import Crew from crewai.crew import Crew
@@ -1798,7 +1799,9 @@ def test_crew_train_success(task_evaluator, crew_training_handler, kickoff):
agents=[researcher, writer], agents=[researcher, writer],
tasks=[task], 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( task_evaluator.assert_has_calls(
[ [
mock.call(researcher), mock.call(researcher),
@@ -1882,7 +1885,7 @@ def test__setup_for_training():
for agent in agents: for agent in agents:
assert agent.allow_delegation is True assert agent.allow_delegation is True
crew._setup_for_training() crew._setup_for_training("trained_agents_data.pkl")
assert crew._train is True assert crew._train is True
assert task.human_input is True assert task.human_input is True