diff --git a/setup_env.py b/setup_env.py index 3bf5fb8f7..4ef7683b8 100644 --- a/setup_env.py +++ b/setup_env.py @@ -104,7 +104,7 @@ def run_command(command, shell=False, log_step=None): subprocess.run(command, shell=shell, check=True) except subprocess.CalledProcessError as e: logging.error(f"Error occurred while running command: {e}") - sys.exit(1) + sys.exit(1) def prepare_model(): _, arch = system_info() diff --git a/tests/test_command_wrappers.py b/tests/test_command_wrappers.py new file mode 100644 index 000000000..2ad3061c1 --- /dev/null +++ b/tests/test_command_wrappers.py @@ -0,0 +1,49 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import setup_env +from utils import e2e_benchmark + + +def test_setup_env_run_command_does_not_exit_on_success_without_log_step(): + with patch("setup_env.subprocess.run") as mock_run, patch("setup_env.sys.exit") as mock_exit: + setup_env.run_command(["echo", "ok"]) + + mock_run.assert_called_once_with(["echo", "ok"], shell=False, check=True) + mock_exit.assert_not_called() + + +def test_setup_env_run_command_exits_on_failure_without_log_step(): + error = setup_env.subprocess.CalledProcessError(1, ["echo", "fail"]) + with patch("setup_env.subprocess.run", side_effect=error), patch( + "setup_env.sys.exit" + ) as mock_exit: + setup_env.run_command(["echo", "fail"]) + + mock_exit.assert_called_once_with(1) + + +def test_e2e_benchmark_run_command_does_not_exit_on_success_without_log_step(): + with patch("utils.e2e_benchmark.subprocess.run") as mock_run, patch( + "utils.e2e_benchmark.sys.exit" + ) as mock_exit: + e2e_benchmark.run_command(["echo", "ok"]) + + mock_run.assert_called_once_with(["echo", "ok"], shell=False, check=True) + mock_exit.assert_not_called() + + +def test_e2e_benchmark_run_benchmark_uses_run_command_without_exiting_early(): + e2e_benchmark.args = SimpleNamespace( + model="model.gguf", + n_token=32, + threads=4, + n_prompt=64, + ) + + with patch("utils.e2e_benchmark.platform.system", return_value="Linux"), patch( + "utils.e2e_benchmark.os.path.exists", return_value=True + ), patch("utils.e2e_benchmark.run_command") as mock_run: + e2e_benchmark.run_benchmark() + + mock_run.assert_called_once() diff --git a/utils/e2e_benchmark.py b/utils/e2e_benchmark.py index 07f93ed72..7eb8211b4 100644 --- a/utils/e2e_benchmark.py +++ b/utils/e2e_benchmark.py @@ -20,7 +20,7 @@ def run_command(command, shell=False, log_step=None): subprocess.run(command, shell=shell, check=True) except subprocess.CalledProcessError as e: logging.error(f"Error occurred while running command: {e}") - sys.exit(1) + sys.exit(1) def run_benchmark(): build_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "build") @@ -56,4 +56,4 @@ def parse_args(): if __name__ == "__main__": logging.basicConfig(level=logging.INFO) args = parse_args() - run_benchmark() \ No newline at end of file + run_benchmark()