Files
crewAI/lib/cli/tests/test_train_crew.py
T
Sathyajith MCandVidit Ostwal e836a191cf
CodeQL Advanced / Analyze (actions) (push) Has been cancelled
CodeQL Advanced / Analyze (python) (push) Has been cancelled
Vulnerability Scan / Detect changes (push) Has been cancelled
Vulnerability Scan / pip-audit (push) Has been cancelled
Nightly Canary Release / Check for new commits (push) Has been cancelled
Nightly Canary Release / Build nightly packages (push) Has been cancelled
Nightly Canary Release / Publish nightly to PyPI (push) Has been cancelled
fix(cli): correct train filename error message (#7903)
* fix(cli): correct train filename error message

* fix(cli): fail train on invalid input with nonzero exit

---------

Co-authored-by: Vidit Ostwal <110953813+Vidit-Ostwal@users.noreply.github.com>
2026-10-06 18:46:20 +00:00

104 lines
3.2 KiB
Python

import subprocess
from unittest import mock
import click
import pytest
from crewai_cli.train_crew import train_crew
@mock.patch("crewai_cli.train_crew.subprocess.run")
def test_train_crew_positive_iterations(mock_subprocess_run):
n_iterations = 5
mock_subprocess_run.return_value = subprocess.CompletedProcess(
args=["uv", "run", "train", str(n_iterations)],
returncode=0,
stdout="Success",
stderr="",
)
train_crew(n_iterations, "trained_agents_data.pkl")
mock_subprocess_run.assert_called_once_with(
["uv", "run", "train", str(n_iterations), "trained_agents_data.pkl"],
capture_output=False,
text=True,
check=True,
)
@mock.patch("crewai_cli.train_crew.subprocess.run")
def test_train_crew_zero_iterations(mock_subprocess_run):
with pytest.raises(
click.ClickException,
match="The number of iterations must be a positive integer.",
):
train_crew(0, "trained_agents_data.pkl")
mock_subprocess_run.assert_not_called()
@mock.patch("crewai_cli.train_crew.subprocess.run")
def test_train_crew_negative_iterations(mock_subprocess_run):
with pytest.raises(
click.ClickException,
match="The number of iterations must be a positive integer.",
):
train_crew(-2, "trained_agents_data.pkl")
mock_subprocess_run.assert_not_called()
@mock.patch("crewai_cli.train_crew.subprocess.run")
def test_train_crew_invalid_filename(mock_subprocess_run):
with pytest.raises(
click.ClickException, match=r"The filename must end with \.pkl"
):
train_crew(5, "myfile.txt")
mock_subprocess_run.assert_not_called()
@mock.patch("crewai_cli.train_crew.click")
@mock.patch("crewai_cli.train_crew.subprocess.run")
def test_train_crew_called_process_error(mock_subprocess_run, click):
n_iterations = 5
mock_subprocess_run.side_effect = subprocess.CalledProcessError(
returncode=1,
cmd=["uv", "run", "train", str(n_iterations)],
output="Error",
stderr="Some error occurred",
)
train_crew(n_iterations, "trained_agents_data.pkl")
mock_subprocess_run.assert_called_once_with(
["uv", "run", "train", str(n_iterations), "trained_agents_data.pkl"],
capture_output=False,
text=True,
check=True,
)
click.echo.assert_has_calls(
[
mock.call.echo(
"An error occurred while training the crew: Command '['uv', 'run', 'train', '5']' returned non-zero exit status 1.",
err=True,
),
mock.call.echo("Error", err=True),
]
)
@mock.patch("crewai_cli.train_crew.click")
@mock.patch("crewai_cli.train_crew.subprocess.run")
def test_train_crew_unexpected_exception(mock_subprocess_run, click):
n_iterations = 5
mock_subprocess_run.side_effect = Exception("Unexpected error")
train_crew(n_iterations, "trained_agents_data.pkl")
mock_subprocess_run.assert_called_once_with(
["uv", "run", "train", str(n_iterations), "trained_agents_data.pkl"],
capture_output=False,
text=True,
check=True,
)
click.echo.assert_called_once_with(
"An unexpected error occurred: Unexpected error", err=True
)