'''
RunPod | CLI | Utils | SSH Command
'''
import unittest
from unittest.mock import patch, MagicMock

import paramiko

from runpod.cli.utils.ssh_cmd import SSHConnection


class TestSSHConnection(unittest.TestCase):
    """ Test the SSHConnection class. """

    def setUp(self):

        self.patch_get_pod_ssh_ip_port = patch('runpod.cli.utils.ssh_cmd.get_pod_ssh_ip_port',
                                               return_value=('127.0.0.1', 22)).start()

        self.patch_find_ssh_key_file = patch('runpod.cli.utils.ssh_cmd.find_ssh_key_file',
                                             return_value='key_file').start()

        self.mock_ssh_client = MagicMock()
        patch_paramiko = patch('runpod.cli.utils.ssh_cmd.paramiko.SSHClient',
                               return_value=self.mock_ssh_client).start()

        self.addCleanup(self.patch_get_pod_ssh_ip_port.stop)
        self.addCleanup(self.patch_find_ssh_key_file.stop)
        self.addCleanup(patch_paramiko.stop)

        self.ssh_connection = SSHConnection('pod_id_mock')

    def test_enter(self):
        ''' Test entering the context manager. '''
        self.assertEqual(
            self.ssh_connection, self.ssh_connection.__enter__())  # pylint: disable=unnecessary-dunder-call

    def test_enter_exception(self):
        """ Test entering the context manager with an exception. """
        self.mock_ssh_client.connect.side_effect = paramiko.SSHException
        with self.assertRaises(SystemExit):
            SSHConnection('pod_id_mock')

    def test_exit(self):
        ''' Test exiting the context manager. '''
        self.ssh_connection.__exit__(None, None, None)

        self.mock_ssh_client.close.assert_called_once()

    def test_run_commands(self):
        ''' Test that run_commands() calls exec_command() on the SSH object. '''
        commands = ['command1', 'command2']

        mock_exec_command = self.mock_ssh_client.exec_command
        mock_exec_command.return_value = (None, ['stdout1', 'stdout2'], ['stderr1', 'stderr2'])
        self.ssh_connection.run_commands(commands)

        assert mock_exec_command.call_count == 2

    def test_put_file(self):
        ''' Test that put_file() calls put() on the SFTP object. '''
        local_path = '/local/file.txt'
        remote_path = '/remote/file.txt'

        mock_sftp = self.mock_ssh_client.open_sftp.return_value.__enter__.return_value
        self.ssh_connection.put_file(local_path, remote_path)

        mock_sftp.put.assert_called_once_with(local_path, remote_path)

    def test_get_file(self):
        ''' Test that get_file() calls get() on the SFTP object. '''
        local_path = '/local/file.txt'
        remote_path = '/remote/file.txt'

        mock_sftp = self.mock_ssh_client.open_sftp.return_value.__enter__.return_value
        self.ssh_connection.get_file(remote_path, local_path)

        mock_sftp.get.assert_called_once_with(remote_path, local_path)

    @patch('subprocess.run')
    def test_launch_terminal(self, mock_subprocess):
        ''' Test that launch_terminal() calls subprocess.run(). '''
        self.ssh_connection.launch_terminal()
        mock_subprocess.assert_called_once()

    @patch('subprocess.run')
    def test_rsync(self, mock_subprocess):
        ''' Test that rsync() calls subprocess.run(). '''
        self.ssh_connection.rsync('local_path', 'remote_path', quiet=True)
        mock_subprocess.assert_called_once()

    # Test that the signal handler closes the connection.
    @patch('runpod.cli.utils.ssh_cmd.SSHConnection.close')
    def test_signal_handler(self, mock_close):
        ''' Test that the signal handler closes the connection. '''
        with patch('sys.exit') as mock_exit:
            self.ssh_connection._signal_handler(None, None)  # pylint: disable=protected-access
        mock_close.assert_called_once()
        assert mock_exit.called
