virtual-deployment/virtualbox/pybox/tests/test_install_vbox.py

import install_vbox
import unittest

from consts.networking import OAM, MGMT
from dataclasses import dataclass
from unittest.mock import MagicMock, patch, call, ANY


# Network
OAM_CONFIG = [getattr(OAM, attr) for attr in dir(OAM) if not attr.startswith('__')]
MGMT_CONFIG = [getattr(MGMT, attr) for attr in dir(MGMT) if not attr.startswith('__')]


@dataclass
class VirtualBoxOptions:
    vboxnet_type: str
    labname: str

    controllers: int
    workers: int
    storages: int

    username: str
    password: str

    force_delete_lab: str

    setup_type: str
    securityprofile: str
    lowlatency: str
    install_mode: str

    controller0_ip: str
    controller1_ip: str
    vboxnet_ip: str

    config_controller_ini: str
    ini_oam_cidr: str
    ini_oam_ip_start_address: str
    ini_oam_ip_end_address: str

    ansible_controller_config: str

    nat_controller_floating_local_ssh_port: str
    nat_controller0_local_ssh_port: str
    nat_controller_1_local_ssh_port: str
    controller_floating_ip: str

    config_files_dir: str
    config_files_dir_dont_follow_links: str
    lab_setup_conf: str

    script1: str
    script2: str
    script3: str
    script4: str
    script5: str

    hostiocache: str

    list_stages: str
    logpath: str
    custom_stages: str
    from_stage: str
    to_stage: str

    snapshot: str


class MenuSelectorTestCase(unittest.TestCase):
    """
    Class to test menu_selector method
    """

    @patch("install_vbox.serial")
    def test_menu_selector(self, mock_serial):
        """
        Test menu_selector method
        """

        # Setup
        mock_stream = MagicMock()
        setup_type_aio_sx = "AIO-SX"
        setup_type_aio_dx = "AIO-DX"
        setup_type_other = "OTHER"
        security_profile_extended = "extended"
        security_profile_other = "other"
        low_latency_true = True
        low_latency_false = False
        install_mode_serial = "serial"
        install_mode_graphical = "graphical"

        permutations = [
            ((setup_type_aio_sx, security_profile_extended, low_latency_true, install_mode_serial), 6),
            ((setup_type_aio_sx, security_profile_other, low_latency_false, install_mode_serial), 4),
            ((setup_type_aio_dx, security_profile_extended, low_latency_true, install_mode_graphical), 7),
            ((setup_type_aio_dx, security_profile_other, low_latency_false, install_mode_graphical), 5),
            ((setup_type_other, security_profile_extended, low_latency_true, install_mode_serial), 4),
            ((setup_type_other, security_profile_other, low_latency_false, install_mode_graphical), 4)
        ]

        # Run
        accumulated_calls = 0
        for permutation, calls_number in permutations:
            with self.subTest(permutation=permutation):
                install_vbox.menu_selector(mock_stream, *permutation)

                # Assert
                accumulated_calls += calls_number
                self.assertEqual(mock_serial.send_bytes.call_count, accumulated_calls)


class SetupNetworkingTestCase(unittest.TestCase):
    """
    Class to test setup_networking method
    """

    mock_stream = MagicMock()
    mock_ip = "192.168.1.1"
    mock_gateway_ip = "192.168.1.254"
    mock_password = "password"

    @patch("install_vbox.serial")
    @patch("install_vbox.LOG")
    @patch("install_vbox.subprocess.call")
    @patch("install_vbox.host_helper")
    def test_setup_networking(
        self,
        mock_host_helper,
        mock_subprocess_call,
        mock_log,
        mock_serial,
    ):
        """
        Test setup_networking
        """

        # Setup
        mock_subprocess_call.return_value = 0
        v_box = VirtualBoxOptions
        v_box.vboxnet_type = "hostonly"
        install_vbox.V_BOX_OPTIONS = v_box

        # Run
        install_vbox.setup_networking(self.mock_stream, self.mock_ip, self.mock_gateway_ip, password=self.mock_password)

        # Assert
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               "/sbin/ip address list",
                                               prompt=self.mock_ip,
                                               fail_ok=True,
                                               timeout=10)
        mock_host_helper.check_password.assert_has_calls([
            call(self.mock_stream, password=self.mock_password),
            call(self.mock_stream, password=self.mock_password),
            call(self.mock_stream, password=self.mock_password)
        ])
        oam_if = OAM_CONFIG[0]['device']
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               f"sudo /sbin/ip addr add {self.mock_ip}/24 dev {oam_if}",
                                               expect_prompt=False)
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               f"sudo /sbin/ip link set {oam_if} up",
                                               expect_prompt=False)
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               f"sudo route add default gw {self.mock_gateway_ip}",
                                               expect_prompt=False)
        self.assertEqual(mock_subprocess_call.call_args_list, [call(['ping', '-c', '1', self.mock_ip])])
        mock_log.info.assert_any_call("Ping succeeded!")


class FixNetworkingTestCase(unittest.TestCase):
    """
    Class to test fix_networking method
    """

    mock_stream = MagicMock()
    mock_release_r2 = "R2"
    mock_release_r3 = "R3"
    mock_password = "Li69nux*"

    @patch("install_vbox.serial")
    @patch("install_vbox.host_helper")
    def test_fix_networking_r2(self, mock_host_helper, mock_serial):
        """
        Test fix_networking for release R2
        """

        # Run
        install_vbox.fix_networking(self.mock_stream, self.mock_release_r2, self.mock_password)

        # Assert
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               "sudo /sbin/ip link set eth0 down",
                                               expect_prompt=False)
        mock_host_helper.check_password.assert_called_with(self.mock_stream, password=self.mock_password)
        mock_serial.send_bytes.assert_any_call(
            self.mock_stream,
            "sudo /sbin/ip link set eth0 up",
            expect_prompt=False)
        mock_host_helper.check_password.assert_called_with(self.mock_stream, password=self.mock_password)

    @patch("install_vbox.serial")
    @patch("install_vbox.host_helper")
    def test_fix_networking_not_r2(self, mock_host_helper, mock_serial):
        """
        Test fix_networking for releases other than R2
        """

        # Run
        install_vbox.fix_networking(self.mock_stream, self.mock_release_r3, self.mock_password)

        # Assert
        oam_if = OAM_CONFIG[0]['device']
        mock_serial.send_bytes.assert_any_call(self.mock_stream,
                                               f"sudo /sbin/ip link set {oam_if} down",
                                               expect_prompt=False)
        mock_host_helper.check_password.assert_called_with(self.mock_stream, password=self.mock_password)
        mock_serial.send_bytes.assert_any_call(
            self.mock_stream,
            f"sudo /sbin/ip link set {oam_if} up",
            expect_prompt=False)
        mock_host_helper.check_password.assert_called_with(self.mock_stream, password=self.mock_password)


class InstallController0TestCase(unittest.TestCase):
    """
    Class to test install_controller_0 method
    """

    mock_stream = MagicMock()
    mock_menu_select_dict = {
        "setup_type": "Duplex",
        "securityprofile": "Standard",
        "lowlatency": False,
        "install_mode": "standard"
    }
    mock_network_dict = {
        "ctrlr0_ip": "192.168.1.2",
        "gateway_ip": "192.168.1.1",
        "username": "wrsroot",
        "password": "Li69nux*"
    }

    @patch("install_vbox.serial")
    @patch("install_vbox.host_helper")
    @patch("install_vbox.menu_selector")
    @patch("install_vbox.setup_networking")
    def test_install_controller_0(
        self, mock_setup_networking, mock_menu_selector, mock_host_helper, mock_serial
    ):
        """
        Test install_controller_0
        """

        # Run
        install_vbox.install_controller_0(self.mock_stream, self.mock_menu_select_dict, self.mock_network_dict)

        # Assert
        mock_menu_selector.assert_called_once_with(
            self.mock_stream,
            self.mock_menu_select_dict["setup_type"],
            self.mock_menu_select_dict["securityprofile"],
            self.mock_menu_select_dict["lowlatency"],
            self.mock_menu_select_dict["install_mode"]
        )
        mock_serial.expect_bytes.assert_called_with(
            self.mock_stream,
            "login:",
            timeout=ANY)
        mock_host_helper.change_password.assert_called_once_with(
            self.mock_stream,
            username=self.mock_network_dict["username"],
            password=self.mock_network_dict["password"]
        )
        mock_host_helper.disable_logout.assert_called_once_with(self.mock_stream)
        mock_setup_networking.assert_called_once_with(
            self.mock_stream,
            self.mock_network_dict["ctrlr0_ip"],
            self.mock_network_dict["gateway_ip"],
            password=self.mock_network_dict["password"]
        )

    @patch("serial.LOG.info")
    @patch("install_vbox.serial")
    @patch("install_vbox.host_helper")
    @patch("install_vbox.time")
    @patch("install_vbox.menu_selector")
    @patch("install_vbox.setup_networking")
    def test_install_controller_0_exception(
        self, mock_setup_networking, mock_menu_selector, mock_time, mock_host_helper, mock_serial, mock_log_info
    ):
        """
        Test install_controller_0 when an exception occurs during login
        """

        # Setup
        mock_serial.expect_bytes.side_effect = [Exception(), None]
        mock_time.time.return_value = 0

        # Run
        install_vbox.install_controller_0(self.mock_stream, self.mock_menu_select_dict, self.mock_network_dict)

        # Assert
        self.assertEqual(mock_serial.expect_bytes.call_count, 2)
        self.assertEqual(mock_serial.expect_bytes.call_args_list[1][1]["timeout"], ANY)
        mock_menu_selector.assert_called_once()
        mock_setup_networking.assert_called_once()
        mock_host_helper.change_password.assert_called_once()
        mock_host_helper.disable_logout.assert_called_once()
        self.assertEqual(mock_log_info.call_count, 4)


class DeleteLabTestCase(unittest.TestCase):
    """
    Class to test delete_lab method
    """

    mock_labname = "test_lab"
    mock_node_list = ["vm1", "vm2", "vm3"]

    @patch("install_vbox.vboxmanage")
    @patch("install_vbox.LOG")
    @patch("install_vbox.time")
    @patch("install_vbox.input", return_value="y")
    def test_delete_lab_not_force(
        self, mock_input, mock_time, mock_log, mock_vboxmanage
    ):
        """
        Test delete_lab with force=False and user input 'y'
        """

        # Setup
        mock_vboxmanage.get_all_vms.return_value = self.mock_node_list

        # Run
        install_vbox.delete_lab(self.mock_labname, force=False)

        # Assert
        mock_vboxmanage.get_all_vms.assert_called_once_with(self.mock_labname, option="vms")
        mock_input.assert_called_once_with()
        mock_log.info.assert_has_calls([
            call("This will delete lab %s with vms: %s", self.mock_labname, self.mock_node_list),
            call("Continue? (y/N)"),
            call("#### Deleting lab %s.", self.mock_labname),
            call("VMs in lab: %s.", self.mock_node_list),
        ])
        mock_vboxmanage.vboxmanage_controlvms.assert_called_once_with(self.mock_node_list, "poweroff")
        mock_time.sleep.assert_called_once_with(2)
        mock_vboxmanage.vboxmanage_deletevms.assert_called_once_with(self.mock_node_list)

    @patch("install_vbox.LOG")
    @patch("install_vbox.vboxmanage")
    @patch("install_vbox.input", return_value="n")
    def test_delete_lab_not_force_abort(
        self, mock_input, mock_vboxmanage, mock_log
    ):
        """
        Test delete_lab with force=False and user input 'n'
        """

        # Setup
        mock_vboxmanage.get_all_vms.return_value = self.mock_node_list

        # Run
        with self.assertRaises(SystemExit):
            install_vbox.delete_lab(self.mock_labname, force=False)

        # Assert
        mock_input.assert_called_once_with()
        mock_log.info.assert_called_with("Aborting!")

    @patch("install_vbox.vboxmanage")
    @patch("install_vbox.LOG")
    @patch("install_vbox.time")
    def test_delete_lab_force(
        self, mock_time, mock_log, mock_vboxmanage
    ):
        """
        Test delete_lab with force=True
        """

        # Setup
        mock_vboxmanage.get_all_vms.return_value = self.mock_node_list

        # Run
        install_vbox.delete_lab(self.mock_labname, force=True)

        # Assert
        mock_vboxmanage.get_all_vms.assert_called_once_with(self.mock_labname, option="vms")
        mock_log.info.assert_has_calls([
            call("#### Deleting lab %s.", self.mock_labname),
            call("VMs in lab: %s.", self.mock_node_list),
        ])
        mock_vboxmanage.vboxmanage_controlvms.assert_called_once_with(self.mock_node_list, "poweroff")
        mock_time.sleep.assert_called_once_with(2)
        mock_vboxmanage.vboxmanage_deletevms.assert_called_once_with(self.mock_node_list)


class GetDiskSizesTestCase(unittest.TestCase):
    """
    Class to test get_disk_sizes method
    """

    def test_get_disk_sizes_valid_input(self):
        """
        Test get_disk_sizes with valid input
        """

        # Setup
        valid_input = "100,200,300"

        # Run
        result = install_vbox.get_disk_sizes(valid_input)

        # Assert
        self.assertEqual(result, ['100', '200', '300'])

    @patch("install_vbox.LOG")
    def test_get_disk_sizes_invalid_input(self, mock_log):
        """
        Test get_disk_sizes with invalid input
        """

        # Setup
        invalid_input = "-100,200,300"

        # Assert
        with self.assertRaises(Exception) as context:
            install_vbox.get_disk_sizes(invalid_input)

        self.assertTrue("Disk sizes must be a comma separated list of positive integers." in str(context.exception))
        mock_log.info.assert_called_once_with("Disk sizes must be a comma separated list of positive integers.")


class TestCreateLab(unittest.TestCase):
    """
    Class to test create_lab method
    """

    # This function needs to be refactored in order to be tested
    pass


class TestGetHostnames(unittest.TestCase):
    """
    Class to test get_hostnames method
    """

    @patch.object(install_vbox, 'V_BOX_OPTIONS', create=True)
    def test_get_hostnames(self, mock_options):
        """
        Test get_hostnames
        """

        # Setup
        mock_options.controllers = 2
        mock_options.workers = 2
        mock_options.storages = 2
        mock_options.labname = "test"

        expected = {
            'test-controller-0': 'controller-0',
            'test-controller-1': 'controller-1',
            'test-worker-0': f'worker-{id}',
            'test-worker-1': f'worker-{id}',
            'test-storage-0': 'storage-0',
            'test-storage-1': 'storage-1',
        }

        # Run
        result = install_vbox.get_hostnames()

        # Assert
        self.assertEqual(result, expected)

    @patch.object(install_vbox, 'V_BOX_OPTIONS', create=True)
    def test_get_hostnames_with_ignore(self, mock_options):
        """
        Test get_hostnames with ignore
        """

        # Setup
        mock_options.controllers = 2
        mock_options.workers = 2
        mock_options.storages = 2
        mock_options.labname = "test"

        ignore = ['test-controller-0', 'test-worker-1']
        expected = {
            'test-controller-1': 'controller-1',
            'test-worker-0': f'worker-{id}',
            'test-storage-0': 'storage-0',
            'test-storage-1': 'storage-1',
        }

        # Run
        result = install_vbox.get_hostnames(ignore=ignore)

        # Assert
        self.assertEqual(result, expected)

    @patch.object(install_vbox, 'V_BOX_OPTIONS', create=True)
    def test_get_hostnames_with_selected_personalities(self, mock_options):
        """
        Test get_hostnames with selected personalities
        """

        # Setup
        mock_options.controllers = 2
        mock_options.workers = 2
        mock_options.storages = 2
        mock_options.labname = "test"

        personalities = ['controller', 'worker']
        expected = {
            'test-controller-0': 'controller-0',
            'test-controller-1': 'controller-1',
            'test-worker-0': f'worker-{id}',
            'test-worker-1': f'worker-{id}',
        }

        # Run
        result = install_vbox.get_hostnames(personalities=personalities)

        # Assert
        self.assertEqual(result, expected)


class TestGetPersonalities(unittest.TestCase):
    """
    Class to test get_personalities method
    """

    @patch.object(install_vbox, 'V_BOX_OPTIONS', create=True)
    def test_get_personalities(self, mock_options):
        """
        Test get_personalities
        """

        # Setup
        mock_options.controllers = 2
        mock_options.workers = 2
        mock_options.storages = 2
        mock_options.labname = "test"

        expected = {
            'test-controller-0': 'controller',
            'test-controller-1': 'controller',
            'test-worker-0': 'worker',
            'test-worker-1': 'worker',
            'test-storage-0': 'storage',
            'test-storage-1': 'storage',
        }

        # Run and Assert
        self.assertEqual(install_vbox.get_personalities(), expected)

    @patch.object(install_vbox, 'V_BOX_OPTIONS', create=True)
    def test_get_personalities_with_ignore(self, mock_options):
        """
        Test get_personalities with ignore
        """

        # Setup
        mock_options.controllers = 2
        mock_options.workers = 2
        mock_options.storages = 2
        mock_options.labname = "test"

        ignore = ['test-controller-0', 'test-worker-1']
        expected = {
            'test-controller-1': 'controller',
            'test-worker-0': 'worker',
            'test-storage-0': 'storage',
            'test-storage-1': 'storage',
        }

        # Run and Assert
        self.assertEqual(install_vbox.get_personalities(ignore=ignore), expected)


class TestCreateHostBulkAdd(unittest.TestCase):
    """
    Class to test create_host_bulk_add method
    """

    @patch.object(install_vbox, 'V_BOX_OPTIONS', create=True)
    @patch.object(install_vbox, 'vboxmanage', create=True)
    @patch.object(install_vbox, 'get_personalities')
    @patch.object(install_vbox, 'get_hostnames')
    def test_create_host_bulk_add(self, mock_get_hostnames, mock_get_personalities, mock_vboxmanage, mock_options):
        """
        Test create_host_bulk_add
        """

        # Setup
        mock_options.labname = "test"
        mock_vboxmanage.get_all_vms.return_value = ['test-controller-0', 'test-controller-1', 'test-worker-0',
                                                    'test-storage-0']
        mock_vboxmanage.vboxmanage_showinfo.return_value = b'macaddress2="080027C95571"\n'
        mock_get_personalities.return_value = {
            'test-controller-1': 'controller',
            'test-worker-0': 'worker',
            'test-storage-0': 'storage',
        }
        mock_get_hostnames.return_value = {
            'test-controller-1': 'controller-1',
            'test-worker-0': 'worker-0',
            'test-storage-0': 'storage-0',
        }
        expected_xml = (
            '<?xml version="1.0" encoding="UTF-8" ?>\n'
            '<hosts>\n'
            '    <host>\n'
            '        <hostname>controller-1</hostname>\n'
            '        <personality>controller</personality>\n'
            '        <mgmt_mac>08:00:27:C9:55:71</mgmt_mac>\n'
            '    </host>\n'
            '    <host>\n'
            '        <hostname>worker-0</hostname>\n'
            '        <personality>worker</personality>\n'
            '        <mgmt_mac>08:00:27:C9:55:71</mgmt_mac>\n'
            '    </host>\n'
            '    <host>\n'
            '        <hostname>storage-0</hostname>\n'
            '        <personality>storage</personality>\n'
            '        <mgmt_mac>08:00:27:C9:55:71</mgmt_mac>\n'
            '    </host>\n'
            '</hosts>\n'
        )

        # Run
        actual_xml = install_vbox.create_host_bulk_add()

        # Assert
        self.assertEqual(actual_xml, expected_xml)


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