virtual-deployment/virtualbox/pybox/utils/tests/test_sftp.py

import unittest
from unittest.mock import MagicMock, patch, call, ANY
import sftp


class SftpSendTestCase(unittest.TestCase):
    """
    Class to test sftp_send method
    """

    @patch("serial.LOG.info")
    @patch("sftp.paramiko.SSHClient")
    def test_sftp_send(self, mock_ssh_client, mock_log_info):
        """
        Test sftp_send method
        """

        # Setup
        source = "/local/path/to/file"
        destination = "/remote/path/to/file"
        client_dict = {
            "remote_host": "127.0.0.1",
            "username": "username",
            "remote_port": 22,
            "password": "password"
        }

        mock_ssh_instance = MagicMock()
        mock_sftp_client_instance = MagicMock()
        mock_ssh_client.return_value = mock_ssh_instance
        mock_ssh_instance.open_sftp.return_value = mock_sftp_client_instance

        # Run
        sftp.sftp_send(source, destination, client_dict)

        # Assert
        mock_ssh_instance.connect.assert_called_once_with(
            client_dict["remote_host"],
            port=client_dict["remote_port"],
            username=client_dict["username"],
            password=client_dict["password"],
            look_for_keys=False,
            allow_agent=False
        )
        mock_sftp_client_instance.put.assert_called_once_with(source, destination)
        mock_sftp_client_instance.close.assert_called_once()
        mock_ssh_instance.close.assert_called_once()
        mock_log_info.assert_any_call(ANY)


class TestSendDir(unittest.TestCase):
    """
    Class to test send_dir method
    """

    @patch("serial.LOG.info")
    @patch('sftp.exec_cmd')
    @patch('sftp.getpass.getuser')
    def test_send_dir(self, mock_getuser, mock_popen, mock_log_info):
        """
        Test send_dir method
        """

        # Setup
        mock_getuser.return_value = 'username'
        process_mock = MagicMock()
        attrs = {'wait.return_value': None,
                 'returncode': 0,
                 'stdout.readline.side_effect': [b'line1\n', b'line2\n', b'']}
        process_mock.configure_mock(**attrs)
        mock_popen.return_value.__enter__.return_value = process_mock

        # Run and assert
        # test without follow_links and clear_known_hosts
        sftp.send_dir({
            'source': 'source',
            'remote_host': 'remote_host',
            'remote_port': 'remote_port',
            'destination': 'destination',
            'username': 'username',
            'password': 'password',
            'follow_links': False,
            'clear_known_hosts': False,
        })
        self.assertEqual(mock_popen.call_count, 1)

        # test with follow_links and clear_known_hosts and localhost
        sftp.send_dir({
            'source': 'source',
            'remote_host': '127.0.0.1',
            'remote_port': 'remote_port',
            'destination': 'destination',
            'username': 'username',
            'password': 'password',
            'follow_links': True,
            'clear_known_hosts': True,
        })
        self.assertEqual(mock_popen.call_count, 3)

        # test with follow_links and clear_known_hosts and non-localhost
        sftp.send_dir({
            'source': 'source',
            'remote_host': 'remote_host',
            'remote_port': 'remote_port',
            'destination': 'destination',
            'username': 'username',
            'password': 'password',
            'follow_links': True,
            'clear_known_hosts': True,
        })
        self.assertEqual(mock_popen.call_count, 5)

        # test with non-zero return code
        process_mock.returncode = 1


class TestSendDirFallback(unittest.TestCase):
    """
        Class to test send_dir_fallback method
    """

    @patch("serial.LOG.info")
    @patch('sftp.os.listdir')
    @patch('sftp.os.path.isfile')
    @patch('sftp.paramiko.SSHClient')
    def test_send_dir_fallback(self, mock_ssh_client, mock_isfile, mock_listdir, mock_log_info):
        """
        Test send_dir_fallback method
        """

        # Setup
        mock_listdir.return_value = ['test1.img', 'test2.txt', 'test3.iso']
        mock_isfile.return_value = True
        mock_sftp_client = MagicMock()
        mock_ssh_client_instance = MagicMock()
        mock_ssh_client_instance.open_sftp.return_value = mock_sftp_client
        mock_ssh_client.return_value = mock_ssh_client_instance

        source = '/path/to/source'
        remote_host = 'remote_host'
        destination = '/path/to/destination'
        username = 'username'
        password = 'password'

        # Run
        sftp.send_dir_fallback(source, remote_host, destination, username, password)

        # Assert
        mock_ssh_client.assert_called_once()
        mock_ssh_client_instance.connect.assert_called_once_with(remote_host, username=username,
                                                                 password=password, look_for_keys=False,
                                                                 allow_agent=False)
        mock_sftp_client.put.assert_any_call(source + 'test1.img', destination + 'images/test1.img')
        mock_sftp_client.put.assert_any_call(source + 'test2.txt', destination + 'test2.txt')
        self.assertNotIn(call(source + 'test3.iso', ANY), mock_sftp_client.put.call_args_list)
        mock_sftp_client.close.assert_called_once()
        mock_ssh_client_instance.close.assert_called_once()
        mock_log_info.assert_any_call(ANY)


if __name__ == '__main__':
    unittest.main()