Fix crew_train_success test

This commit is contained in:
Thiago Moretto
2024-11-13 10:47:49 -03:00
parent 7fb1289205
commit c57cbd8591

View File

@@ -1779,26 +1779,22 @@ def test_crew_train_success(
] ]
) )
crew_training_handler.assert_has_calls( crew_training_handler.assert_any_call("training_data.pkl")
[ crew_training_handler().load.assert_called()
mock.call("training_data.pkl"),
mock.call().load(), crew_training_handler.assert_any_call("trained_agents_data.pkl")
mock.call("trained_agents_data.pkl"), crew_training_handler().load.assert_called()
mock.call().save_trained_data(
agent_id="Researcher", crew_training_handler().save_trained_data.assert_has_calls([
trained_data=task_evaluator().evaluate_training_data().model_dump(), mock.call(
), agent_id="Researcher",
mock.call("trained_agents_data.pkl"), trained_data=task_evaluator().evaluate_training_data().model_dump(),
mock.call().save_trained_data( ),
agent_id="Senior Writer", mock.call(
trained_data=task_evaluator().evaluate_training_data().model_dump(), agent_id="Senior Writer",
), trained_data=task_evaluator().evaluate_training_data().model_dump(),
mock.call(), )
mock.call().load(), ])
mock.call(),
mock.call().load(),
]
)
def test_crew_train_error(): def test_crew_train_error():