|
| 1 | +from csle_common.util.ssh_util import SSHUtil |
| 2 | +from unittest.mock import patch, MagicMock |
| 3 | +import pytest |
| 4 | + |
| 5 | + |
| 6 | +class TestSSHUtilSuite: |
| 7 | + """ |
| 8 | + Test suite for ssh_util |
| 9 | + """ |
| 10 | + |
| 11 | + @pytest.fixture(autouse=True) |
| 12 | + def mock_sleep(self): |
| 13 | + """ |
| 14 | + Mock time.sleep to avoid delays |
| 15 | + """ |
| 16 | + with patch("time.sleep", return_value=None): |
| 17 | + yield |
| 18 | + |
| 19 | + @patch("csle_common.util.ssh_util.SSHUtil.execute_ssh_cmd") |
| 20 | + def test_execute_ssh_cmds(self, mock_execute_ssh_cmd) -> None: |
| 21 | + """ |
| 22 | + Test the method that executes a list of commands over an ssh connection to the emulation |
| 23 | +
|
| 24 | + :param mock_execute_ssh_cmd: mock_execute_ssh_cmd |
| 25 | +
|
| 26 | + :return: None |
| 27 | + """ |
| 28 | + mock_execute_ssh_cmd.return_value = (b"output", b"error", 1.0) |
| 29 | + cmds = ["ls", "pwd", "whoami"] |
| 30 | + conn = MagicMock() |
| 31 | + results = SSHUtil.execute_ssh_cmds(cmds, conn) |
| 32 | + mock_execute_ssh_cmd.assert_called() |
| 33 | + assert results == [(b"output", b"error", 1.0)] * len(cmds) |
| 34 | + |
| 35 | + def test_execute_ssh_cmd(self) -> None: |
| 36 | + """ |
| 37 | + Test the method that executes an action on the emulation over a ssh connection |
| 38 | +
|
| 39 | + :return: None |
| 40 | + """ |
| 41 | + conn = MagicMock() |
| 42 | + mock_transport = MagicMock() |
| 43 | + mock_session = MagicMock() |
| 44 | + mock_session.exit_status_ready.return_value = True |
| 45 | + mock_session.recv_ready.return_value = True |
| 46 | + mock_session.recv_stderr_ready.return_value = True |
| 47 | + mock_session.recv.side_effect = [b"output", b""] |
| 48 | + mock_session.recv_stderr.side_effect = [b"error", b""] |
| 49 | + conn.get_transport.return_value = mock_transport |
| 50 | + mock_transport.open_session.return_value = mock_session |
| 51 | + |
| 52 | + with pytest.raises(ConnectionError, match="Connection failed"): |
| 53 | + SSHUtil.execute_ssh_cmd("ls", conn) |
0 commit comments