virtual-deployment/virtualbox/pybox/helper/tests/test_host_helper.py

"""
Unit tests related to host_helper 
"""

import unittest
from unittest.mock import MagicMock, patch, call
import streamexpect

import host_helper


class UnlockHostTestCase(unittest.TestCase):
    """
    Class to test unlock_host method
    """

    mock_stream = MagicMock()
    mock_hostname = "hostname"

    @patch("host_helper.serial")
    def test_unlock_host_when_locked(self, mock_serial):
        """
        Test unlock_host when locked
        """

        # Run
        result = host_helper.unlock_host(self.mock_stream, self.mock_hostname)

        # Assert
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               f"system host-list | grep {self.mock_hostname}",
                                               expect_prompt=False)
        mock_serial.expect_bytes.assert_called_once_with(self.mock_stream, "locked")
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               f"system host-unlock {self.mock_hostname}",
                                               expect_prompt=False)
        self.assertIsNone(result)
        self.assertEqual(mock_serial.send_bytes.call_count, 2)

    @patch("host_helper.serial")
    def test_unlock_host_when_not_locked(self, mock_serial):
        """
        Test unlock_host when not locked
        """

        # Setup
        mock_serial.expect_bytes.side_effect = streamexpect.ExpectTimeout

        # Run
        result = host_helper.unlock_host(self.mock_stream, self.mock_hostname)

        # Assert
        mock_serial.send_bytes.assert_called_once_with(self.mock_stream,
                                                       f"system host-list | grep {self.mock_hostname}",
                                                       expect_prompt=False)
        mock_serial.expect_bytes.assert_called_once_with(self.mock_stream, "locked")
        self.assertEqual(result, 1)


class LockHostTestCase(unittest.TestCase):
    """
    Class to test lock_host method
    """

    mock_stream = MagicMock()
    mock_hostname = "hostname"

    @patch("host_helper.serial")
    def test_lock_host_when_unlocked(self, mock_serial):
        """
        Test lock_host when host is unlocked
        """

        # Run
        result = host_helper.lock_host(self.mock_stream, self.mock_hostname)

        # Assert
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               f"system host-list |grep {self.mock_hostname}",
                                               expect_prompt=False)
        mock_serial.expect_bytes.assert_called_once_with(self.mock_stream, "unlocked")
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               f"system host-lock {self.mock_hostname}",
                                               expect_prompt="keystone")
        self.assertEqual(mock_serial.send_bytes.call_count, 2)
        self.assertIsNone(result)

    @patch("host_helper.serial")
    def test_lock_host_when_not_unlocked(self, mock_serial):
        """
        Test lock_host when host is not unlocked
        """

        # Setup
        mock_serial.expect_bytes.side_effect = streamexpect.ExpectTimeout

        # Run
        result = host_helper.lock_host(self.mock_stream, self.mock_hostname)

        # Assert
        mock_serial.send_bytes.assert_called_once_with(self.mock_stream,
                                                       f"system host-list |grep {self.mock_hostname}",
                                                       expect_prompt=False)
        mock_serial.expect_bytes.assert_called_once_with(self.mock_stream, "unlocked")
        self.assertEqual(result, 1)


class RebootHostTestCase(unittest.TestCase):
    """
    Class to test reboot_host method
    """

    @patch("host_helper.serial")
    @patch("host_helper.HostTimeout")
    def test_reboot_host(self, mock_host_timeout, mock_serial):
        """
        Test reboot_host method
        """

        # Setup
        mock_stream = MagicMock()
        mock_hostname = "hostname"

        # Run
        host_helper.reboot_host(mock_stream, mock_hostname)

        # Assert
        mock_serial.send_bytes.assert_called_once_with(mock_stream,
                                                       f"system host-reboot {mock_hostname}",
                                                       expect_prompt=False)
        mock_serial.expect_bytes.assert_called_once_with(mock_stream, "rebooting", mock_host_timeout.REBOOT)


class InstallHostTestCase(unittest.TestCase):
    """
    Class to test install_host method
    """

    mock_stream = MagicMock()
    mock_hostname = "hostname"
    mock_host_id = 1

    @patch("host_helper.serial")
    def test_install_host_controller(self, mock_serial):
        """
        Test install_host for controller type host
        """

        # Setup
        mock_host_type = "controller"

        # Run
        host_helper.install_host(self.mock_stream, self.mock_hostname, mock_host_type, self.mock_host_id)

        # Assert
        mock_serial.send_bytes.assert_called_once_with(self.mock_stream,
                                                       f"system host-update {self.mock_host_id} personality=controller",
                                                       expect_prompt=False)

    @patch("host_helper.serial")
    def test_install_host_storage(self, mock_serial):
        """
        Test install_host for storage type host
        """

        # Setup
        mock_host_type = "storage"

        # Run
        host_helper.install_host(self.mock_stream, self.mock_hostname, mock_host_type, self.mock_host_id)

        # Assert
        mock_serial.send_bytes.assert_called_once_with(self.mock_stream,
                                                       f"system host-update {self.mock_host_id} personality=storage",
                                                       expect_prompt=False)

    @patch("host_helper.serial")
    def test_install_host_worker(self, mock_serial):
        """
        Test install_host for worker type host
        """

        # Setup
        mock_host_type = "worker"

        # Run
        host_helper.install_host(self.mock_stream, self.mock_hostname, mock_host_type, self.mock_host_id)

        # Assert
        mock_serial.send_bytes.assert_called_once_with(
            self.mock_stream,
            f"system host-update {self.mock_host_id} personality=worker hostname={self.mock_hostname}",
            expect_prompt=False)


class DisableLogoutTestCase(unittest.TestCase):
    """
    Class to test disable_logout method
    """

    @patch("host_helper.serial")
    def test_disable_logout(self, mock_serial):
        """
        Test disable_logout method
        """

        # Setup
        mock_stream = MagicMock()

        # Run
        host_helper.disable_logout(mock_stream)

        # Assert
        mock_serial.send_bytes.assert_called_once_with(mock_stream, "export TMOUT=0")


class ChangePasswordTestCase(unittest.TestCase):
    """
    Class to test change_password method
    """

    @patch("host_helper.serial")
    def test_change_password(self, mock_serial):
        """
        Test change_password method
        """

        # Setup
        mock_stream = MagicMock()
        mock_username = "sysadmin"
        mock_password = "Li69nux*"

        # Run
        host_helper.change_password(mock_stream, mock_username, mock_password)

        # Assert
        calls = [
            call.send_bytes(mock_stream, mock_username, expect_prompt=False),
            call.expect_bytes(mock_stream, "Password:"),
            call.send_bytes(mock_stream, mock_username, expect_prompt=False),
            call.expect_bytes(mock_stream, "Current password:"),
            call.send_bytes(mock_stream, mock_username, expect_prompt=False),
            call.expect_bytes(mock_stream, "New password:"),
            call.send_bytes(mock_stream, mock_password, expect_prompt=False),
            call.expect_bytes(mock_stream, "Retype new password"),
            call.send_bytes(mock_stream, mock_password)
        ]

        mock_serial.assert_has_calls(calls, any_order=False)


class CheckPasswordTestCase(unittest.TestCase):
    """
    Class to test check_password method
    """

    mock_stream = MagicMock()

    @patch("host_helper.serial")
    def test_check_password_prompt_found(self, mock_serial):
        """
        Test check_password method when password prompt is found
        """

        # Setup
        mock_password = "Li69nux*"
        mock_serial.expect_bytes.return_value = 0

        # Run
        host_helper.check_password(self.mock_stream, mock_password)

        # Assert
        mock_serial.expect_bytes.assert_called_once_with(self.mock_stream, 'assword', fail_ok=True, timeout=5)
        mock_serial.send_bytes.assert_called_once_with(self.mock_stream, mock_password, expect_prompt=False)

    @patch("host_helper.serial")
    def test_check_password_prompt_not_found(self, mock_serial):
        """
        Test check_password method when password prompt is not found
        """

        # Setup
        mock_serial.expect_bytes.return_value = 1

        # Run
        host_helper.check_password(self.mock_stream, "")

        # Assert
        mock_serial.expect_bytes.assert_called_once_with(self.mock_stream, 'assword', fail_ok=True, timeout=5)
        mock_serial.send_bytes.assert_not_called()


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