diff options
Diffstat (limited to 'tests/aiet/test_utils_proc.py')
| -rw-r--r-- | tests/aiet/test_utils_proc.py | 272 |
1 files changed, 272 insertions, 0 deletions
diff --git a/tests/aiet/test_utils_proc.py b/tests/aiet/test_utils_proc.py new file mode 100644 index 0000000..9fb48dd --- /dev/null +++ b/tests/aiet/test_utils_proc.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates. +# SPDX-License-Identifier: Apache-2.0 +# pylint: disable=attribute-defined-outside-init,no-self-use,not-callable +"""Pytests for testing aiet/utils/proc.py.""" +from pathlib import Path +from typing import Any +from unittest import mock + +import psutil +import pytest +from sh import ErrorReturnCode + +from aiet.utils.proc import Command +from aiet.utils.proc import CommandFailedException +from aiet.utils.proc import CommandNotFound +from aiet.utils.proc import parse_command +from aiet.utils.proc import print_command_stdout +from aiet.utils.proc import run_and_wait +from aiet.utils.proc import save_process_info +from aiet.utils.proc import ShellCommand +from aiet.utils.proc import terminate_command +from aiet.utils.proc import terminate_external_process + + +class TestShellCommand: + """Sample class for collecting tests.""" + + def test_shellcommand_default_value(self) -> None: + """Test the instantiation of the class ShellCommand with no parameter.""" + shell_command = ShellCommand() + assert shell_command.base_log_path == "/tmp" + + @pytest.mark.parametrize( + "base_log_path,expected", [("/test", "/test"), ("/asd", "/asd")] + ) + def test_shellcommand_with_param(self, base_log_path: str, expected: str) -> None: + """Test init ShellCommand with different parameters.""" + shell_command = ShellCommand(base_log_path) + assert shell_command.base_log_path == expected + + def test_run_ls(self, monkeypatch: Any) -> None: + """Test a simple ls command.""" + mock_command = mock.MagicMock() + monkeypatch.setattr(Command, "bake", mock_command) + + mock_get_stdout_stderr_paths = mock.MagicMock() + mock_get_stdout_stderr_paths.return_value = ("/tmp/std.out", "/tmp/std.err") + monkeypatch.setattr( + ShellCommand, "get_stdout_stderr_paths", mock_get_stdout_stderr_paths + ) + + shell_command = ShellCommand() + shell_command.run("ls", "-l") + assert mock_command.mock_calls[0] == mock.call(("-l",)) + assert mock_command.mock_calls[1] == mock.call()( + _bg=True, _err="/tmp/std.err", _out="/tmp/std.out", _tee=True, _bg_exc=False + ) + + def test_run_command_not_found(self) -> None: + """Test whe the command doesn't exist.""" + shell_command = ShellCommand() + with pytest.raises(CommandNotFound): + shell_command.run("lsl", "-l") + + def test_get_stdout_stderr_paths_valid_path(self) -> None: + """Test the method to get files to store stdout and stderr.""" + valid_path = "/tmp" + shell_command = ShellCommand(valid_path) + out, err = shell_command.get_stdout_stderr_paths(valid_path, "cmd") + assert out.exists() and out.is_file() + assert err.exists() and err.is_file() + assert "cmd" in out.name + assert "cmd" in err.name + + def test_get_stdout_stderr_paths_not_invalid_path(self) -> None: + """Test the method to get output files with an invalid path.""" + invalid_path = "/invalid/foo/bar" + shell_command = ShellCommand(invalid_path) + with pytest.raises(FileNotFoundError): + shell_command.get_stdout_stderr_paths(invalid_path, "cmd") + + +@mock.patch("builtins.print") +def test_print_command_stdout_alive(mock_print: Any) -> None: + """Test the print command stdout with an alive (running) process.""" + mock_command = mock.MagicMock() + mock_command.is_alive.return_value = True + mock_command.next.side_effect = ["test1", "test2", StopIteration] + + print_command_stdout(mock_command) + + mock_command.assert_has_calls( + [mock.call.is_alive(), mock.call.next(), mock.call.next()] + ) + mock_print.assert_has_calls( + [mock.call("test1", end=""), mock.call("test2", end="")] + ) + + +@mock.patch("builtins.print") +def test_print_command_stdout_not_alive(mock_print: Any) -> None: + """Test the print command stdout with a not alive (exited) process.""" + mock_command = mock.MagicMock() + mock_command.is_alive.return_value = False + mock_command.stdout = "test" + + print_command_stdout(mock_command) + mock_command.assert_has_calls([mock.call.is_alive()]) + mock_print.assert_called_once_with("test") + + +def test_terminate_external_process_no_process(capsys: Any) -> None: + """Test that non existed process could be terminated.""" + mock_command = mock.MagicMock() + mock_command.terminate.side_effect = psutil.Error("Error!") + + terminate_external_process(mock_command) + captured = capsys.readouterr() + assert captured.out == "Unable to terminate process\n" + + +def test_terminate_external_process_case1() -> None: + """Test when process terminated immediately.""" + mock_command = mock.MagicMock() + mock_command.is_running.return_value = False + + terminate_external_process(mock_command) + mock_command.terminate.assert_called_once() + mock_command.is_running.assert_called_once() + + +def test_terminate_external_process_case2() -> None: + """Test when process termination takes time.""" + mock_command = mock.MagicMock() + mock_command.is_running.side_effect = [True, True, False] + + terminate_external_process(mock_command) + mock_command.terminate.assert_called_once() + assert mock_command.is_running.call_count == 3 + + +def test_terminate_external_process_case3() -> None: + """Test when process termination takes more time.""" + mock_command = mock.MagicMock() + mock_command.is_running.side_effect = [True, True, True] + + terminate_external_process( + mock_command, number_of_attempts=2, wait_period=0.1, wait_for_termination=0.1 + ) + assert mock_command.is_running.call_count == 3 + assert mock_command.terminate.call_count == 2 + + +def test_terminate_external_process_case4() -> None: + """Test when process termination takes more time.""" + mock_command = mock.MagicMock() + mock_command.is_running.side_effect = [True, True, False] + + terminate_external_process( + mock_command, number_of_attempts=2, wait_period=0.1, wait_for_termination=0.1 + ) + mock_command.terminate.assert_called_once() + assert mock_command.is_running.call_count == 3 + assert mock_command.terminate.call_count == 1 + + +def test_terminate_command_no_process() -> None: + """Test command termination when process does not exist.""" + mock_command = mock.MagicMock() + mock_command.process.signal_group.side_effect = ProcessLookupError() + + terminate_command(mock_command) + mock_command.process.signal_group.assert_called_once() + mock_command.is_alive.assert_not_called() + + +def test_terminate_command() -> None: + """Test command termination.""" + mock_command = mock.MagicMock() + mock_command.is_alive.return_value = False + + terminate_command(mock_command) + mock_command.process.signal_group.assert_called_once() + + +def test_terminate_command_case1() -> None: + """Test command termination when it takes time..""" + mock_command = mock.MagicMock() + mock_command.is_alive.side_effect = [True, True, False] + + terminate_command(mock_command, wait_period=0.1) + mock_command.process.signal_group.assert_called_once() + assert mock_command.is_alive.call_count == 3 + + +def test_terminate_command_case2() -> None: + """Test command termination when it takes much time..""" + mock_command = mock.MagicMock() + mock_command.is_alive.side_effect = [True, True, True] + + terminate_command(mock_command, number_of_attempts=3, wait_period=0.1) + assert mock_command.is_alive.call_count == 3 + assert mock_command.process.signal_group.call_count == 2 + + +class TestRunAndWait: + """Test run_and_wait function.""" + + @pytest.fixture(autouse=True) + def setup_method(self, monkeypatch: Any) -> None: + """Init test method.""" + self.execute_command_mock = mock.MagicMock() + monkeypatch.setattr( + "aiet.utils.proc.execute_command", self.execute_command_mock + ) + + self.terminate_command_mock = mock.MagicMock() + monkeypatch.setattr( + "aiet.utils.proc.terminate_command", self.terminate_command_mock + ) + + def test_if_execute_command_raises_exception(self) -> None: + """Test if execute_command fails.""" + self.execute_command_mock.side_effect = Exception("Error!") + with pytest.raises(Exception, match="Error!"): + run_and_wait("command", Path.cwd()) + + def test_if_command_finishes_with_error(self) -> None: + """Test if command finishes with error.""" + cmd_mock = mock.MagicMock() + self.execute_command_mock.return_value = cmd_mock + exit_code_mock = mock.PropertyMock( + side_effect=ErrorReturnCode("cmd", bytearray(), bytearray()) + ) + type(cmd_mock).exit_code = exit_code_mock + + with pytest.raises(CommandFailedException): + run_and_wait("command", Path.cwd()) + + @pytest.mark.parametrize("terminate_on_error, call_count", ((False, 0), (True, 1))) + def test_if_command_finishes_with_exception( + self, terminate_on_error: bool, call_count: int + ) -> None: + """Test if command finishes with error.""" + cmd_mock = mock.MagicMock() + self.execute_command_mock.return_value = cmd_mock + exit_code_mock = mock.PropertyMock(side_effect=Exception("Error!")) + type(cmd_mock).exit_code = exit_code_mock + + with pytest.raises(Exception, match="Error!"): + run_and_wait("command", Path.cwd(), terminate_on_error=terminate_on_error) + + assert self.terminate_command_mock.call_count == call_count + + +def test_save_process_info_no_process(monkeypatch: Any, tmpdir: Any) -> None: + """Test save_process_info function.""" + mock_process = mock.MagicMock() + monkeypatch.setattr("psutil.Process", mock.MagicMock(return_value=mock_process)) + mock_process.children.side_effect = psutil.NoSuchProcess(555) + + pid_file_path = Path(tmpdir) / "test.pid" + save_process_info(555, pid_file_path) + assert not pid_file_path.exists() + + +def test_parse_command() -> None: + """Test parse_command function.""" + assert parse_command("1.sh") == ["bash", "1.sh"] + assert parse_command("1.sh", shell="sh") == ["sh", "1.sh"] + assert parse_command("command") == ["command"] + assert parse_command("command 123 --param=1") == ["command", "123", "--param=1"] |