From 2520f389f214732d8f1260a36d7a813898517a8d Mon Sep 17 00:00:00 2001 From: Eduardo Chiarotti Date: Thu, 16 May 2024 19:43:42 -0300 Subject: [PATCH] feat: add the tests --- tests/cli/cli_test.py | 59 ++++++++++++++++++++++++ tests/cli/train_crew_test.py | 87 ++++++++++++++++++++++++++++++++++++ tests/cli_test.py | 20 --------- tests/crew_test.py | 32 +++++++++++++ 4 files changed, 178 insertions(+), 20 deletions(-) create mode 100644 tests/cli/cli_test.py create mode 100644 tests/cli/train_crew_test.py delete mode 100644 tests/cli_test.py diff --git a/tests/cli/cli_test.py b/tests/cli/cli_test.py new file mode 100644 index 000000000..f2c6879b0 --- /dev/null +++ b/tests/cli/cli_test.py @@ -0,0 +1,59 @@ +from unittest import mock + +import pytest +from click.testing import CliRunner + +from crewai.cli.cli import train, version + + +@pytest.fixture +def runner(): + return CliRunner() + + +@mock.patch("crewai.cli.cli.train_crew") +def test_train_default_iterations(train_crew, runner): + result = runner.invoke(train) + + train_crew.assert_called_once_with(5) + assert result.exit_code == 0 + assert "Training the crew for 5 iterations" in result.output + + +@mock.patch("crewai.cli.cli.train_crew") +def test_train_custom_iterations(train_crew, runner): + result = runner.invoke(train, ["--n_iterations", "10"]) + + train_crew.assert_called_once_with(10) + assert result.exit_code == 0 + assert "Training the crew for 10 iterations" in result.output + + +@mock.patch("crewai.cli.cli.train_crew") +def test_train_invalid_string_iterations(train_crew, runner): + result = runner.invoke(train, ["--n_iterations", "invalid"]) + + train_crew.assert_not_called() + assert result.exit_code == 2 + assert ( + "Usage: train [OPTIONS]\nTry 'train --help' for help.\n\nError: Invalid value for '-n' / '--n_iterations': 'invalid' is not a valid integer.\n" + in result.output + ) + + +def test_version_command(runner): + result = runner.invoke(version) + + assert result.exit_code == 0 + assert "crewai version:" in result.output + + +def test_version_command_with_tools(runner): + result = runner.invoke(version, ["--tools"]) + + assert result.exit_code == 0 + assert "crewai version:" in result.output + assert ( + "crewai tools version:" in result.output + or "crewai tools not installed" in result.output + ) diff --git a/tests/cli/train_crew_test.py b/tests/cli/train_crew_test.py new file mode 100644 index 000000000..9d0d3d4a7 --- /dev/null +++ b/tests/cli/train_crew_test.py @@ -0,0 +1,87 @@ +import subprocess +from unittest import mock + +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): + # Arrange + n_iterations = 5 + mock_subprocess_run.return_value = subprocess.CompletedProcess( + args=["poetry", "run", "train", str(n_iterations)], + returncode=0, + stdout="Success", + stderr="", + ) + + # Act + train_crew(n_iterations) + + # Assert + mock_subprocess_run.assert_called_once_with( + ["poetry", "run", "train", str(n_iterations)], + capture_output=False, + text=True, + check=True, + ) + + +@mock.patch("crewai.cli.train_crew.click") +def test_train_crew_zero_iterations(click): + train_crew(0) + click.echo.assert_called_once_with( + "An unexpected error occurred: The number of iterations must be a positive integer.", + err=True, + ) + + +@mock.patch("crewai.cli.train_crew.click") +def test_train_crew_negative_iterations(click): + train_crew(-2) + click.echo.assert_called_once_with( + "An unexpected error occurred: The number of iterations must be a positive integer.", + err=True, + ) + + +@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=["poetry", "run", "train", str(n_iterations)], + output="Error", + stderr="Some error occurred", + ) + train_crew(n_iterations) + + mock_subprocess_run.assert_called_once_with( + ["poetry", "run", "train", "5"], capture_output=False, text=True, check=True + ) + click.echo.assert_has_calls( + [ + mock.call.echo( + "An error occurred while training the crew: Command '['poetry', '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): + # Arrange + n_iterations = 5 + mock_subprocess_run.side_effect = Exception("Unexpected error") + train_crew(n_iterations) + + mock_subprocess_run.assert_called_once_with( + ["poetry", "run", "train", "5"], capture_output=False, text=True, check=True + ) + click.echo.assert_called_once_with( + "An unexpected error occurred: Unexpected error", err=True + ) diff --git a/tests/cli_test.py b/tests/cli_test.py deleted file mode 100644 index 88247ec69..000000000 --- a/tests/cli_test.py +++ /dev/null @@ -1,20 +0,0 @@ -from click.testing import CliRunner -from crewai.cli.cli import version - - -def test_version_command(): - runner = CliRunner() - result = runner.invoke(version) - assert result.exit_code == 0 - assert "crewai version:" in result.output - - -def test_version_command_with_tools(): - runner = CliRunner() - result = runner.invoke(version, ["--tools"]) - assert result.exit_code == 0 - assert "crewai version:" in result.output - assert ( - "crewai tools version:" in result.output - or "crewai tools not installed" in result.output - ) diff --git a/tests/crew_test.py b/tests/crew_test.py index b9923e2ea..75868ac9e 100644 --- a/tests/crew_test.py +++ b/tests/crew_test.py @@ -993,3 +993,35 @@ def test_manager_agent_with_tools_raises_exception(): with pytest.raises(Exception): crew.kickoff() + + +def test_crew_train_success(): + 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.", + ) + + crew = Crew( + agents=[researcher, writer], + tasks=[task], + ) + + crew.train(n_iterations=2) + + +def test_crew_train_error(): + task = Task( + description="Come up with a list of 5 interesting ideas to explore for an article", + expected_output="5 bullet points with a paragraph for each idea.", + ) + + crew = Crew( + agents=[researcher, writer], + tasks=[task], + ) + + with pytest.raises(TypeError) as e: + crew.train() + assert "train() missing 1 required positional argument: 'n_iterations'" in str( + e + )