Exemple #1
0
    def test_batchai_update_nodes_information(self):
        """Test updating of nodes information."""
        params = ClusterCreateParameters(
            location='eastus',
            vm_size='STANDARD_D1',
            user_account_settings=UserAccountSettings(
                admin_user_name='name', admin_user_password='******'))
        # Update to autoscale Ubuntu DSVM.
        update_nodes_information(params, 'ubuntudsvm', 'Standard_NC6', 1, 3)
        self.assertEquals(params.vm_size, 'Standard_NC6')
        self.assertEquals(
            params.virtual_machine_configuration.image_reference,
            ImageReference('microsoft-ads', 'linux-data-science-vm-ubuntu',
                           'linuxdsvmubuntu'))
        self.assertEquals(params.scale_settings,
                          ScaleSettings(auto_scale=AutoScaleSettings(1, 3)))

        # Update to manual scale Ubuntu LTS.
        update_nodes_information(params, 'UbuntuLTS', 'Standard_NC6', 2, 2)
        self.assertEquals(params.vm_size, 'Standard_NC6')
        self.assertEquals(
            params.virtual_machine_configuration.image_reference,
            ImageReference('Canonical', 'UbuntuServer', '16.04-LTS'))
        self.assertEquals(params.scale_settings,
                          ScaleSettings(manual=ManualScaleSettings(2)))

        # Update image.
        update_nodes_information(params, 'UbuntuDsvm', None, 0, None)
        self.assertEquals(params.vm_size, 'Standard_NC6')
        self.assertEquals(
            params.virtual_machine_configuration.image_reference,
            ImageReference('microsoft-ads', 'linux-data-science-vm-ubuntu',
                           'linuxdsvmubuntu'))
        self.assertEquals(params.scale_settings,
                          ScaleSettings(manual=ManualScaleSettings(2)))

        # Update nothing.
        update_nodes_information(params, None, None, 0, None)
        self.assertEquals(params.vm_size, 'Standard_NC6')
        self.assertEquals(
            params.virtual_machine_configuration.image_reference,
            ImageReference('microsoft-ads', 'linux-data-science-vm-ubuntu',
                           'linuxdsvmubuntu'))
        self.assertEquals(params.scale_settings,
                          ScaleSettings(manual=ManualScaleSettings(2)))

        # Wrong image.
        with self.assertRaises(CLIError):
            update_nodes_information(params, 'unsupported', None, 0, None)

        # No VM size.
        params.vm_size = None
        with self.assertRaises(CLIError):
            update_nodes_information(params, 'unsupported', None, 0, None)

        # No scale settings.
        params.vm_size = 'Standard_NC6'
        params.scale_settings = None
        with self.assertRaises(CLIError):
            update_nodes_information(params, None, None, 0, None)
 def test_setup_task_added_to_cluster_without_setup_task(self):
     actual = _add_setup_task('echo hi', '/tmp', ClusterCreateParameters(node_setup=NodeSetup()))
     self.assertEqual(actual,
                      ClusterCreateParameters(
                          node_setup=NodeSetup(
                              setup_task=SetupTask(
                                  command_line='echo hi',
                                  std_out_err_path_prefix='/tmp',
                                  run_elevated=False
                              ))))
 def test_batchai_update_cluster_user_account_settings_using_command_line_when_key_is_path(self):
     params = ClusterCreateParameters()
     actual = _update_user_account_settings(params, 'user', _data_file('key.txt'), None)
     self.assertEquals(actual.user_account_settings,
                       UserAccountSettings(
                           admin_user_name='user',
                           admin_user_ssh_public_key=SSH_KEY))
 def test_batchai_update_cluster_user_account_settings_using_config(self):
     params = ClusterCreateParameters(
         user_account_settings=UserAccountSettings(
             admin_user_name='name',
             admin_user_password='******'))
     actual = _update_user_account_settings(params, None, None, None)
     self.assertEquals(params, actual)
 def test_batchai_update_cluster_user_account_settings_using_command_line(self):
     params = ClusterCreateParameters()
     actual = _update_user_account_settings(params, 'name', SSH_KEY, 'password')
     self.assertEquals(actual.user_account_settings,
                       UserAccountSettings(
                           admin_user_name='name',
                           admin_user_ssh_public_key=SSH_KEY,
                           admin_user_password='******'))
 def test_setup_task_overwrite(self):
     # Should overwrite setup task but keep other parameters of node setup (here we use mount_volumes to check it)
     actual = _add_setup_task('echo hi', '/tmp',
                              ClusterCreateParameters(
                                  node_setup=NodeSetup(
                                      setup_task=SetupTask(command_line='different'),
                                      mount_volumes=[]
                                  )))
     self.assertEqual(actual,
                      ClusterCreateParameters(
                          node_setup=NodeSetup(
                              setup_task=SetupTask(
                                  command_line='echo hi',
                                  std_out_err_path_prefix='/tmp',
                                  run_elevated=False
                              ),
                              mount_volumes=[])))
 def test_batchai_update_cluster_user_account_settings_using_current_user(self, getuser):
     getuser.return_value = 'current_user'
     params = ClusterCreateParameters()
     actual = _update_user_account_settings(params, None, None, 'password')
     self.assertEquals(actual.user_account_settings,
                       UserAccountSettings(
                           admin_user_name='current_user',
                           admin_user_password='******'))
Exemple #8
0
    def test_batchai_add_azure_file_share_to_cluster_create_parameters(self):
        """Test adding of azure file share into cluster create parameters."""
        params = ClusterCreateParameters(
            location='eastus',
            vm_size='STANDARD_D1',
            user_account_settings=UserAccountSettings(
                admin_user_name='name', admin_user_password='******'))

        # No environment variables given.
        # No environment variables provided.
        os.environ.pop('AZURE_BATCHAI_STORAGE_ACCOUNT', None)
        os.environ.pop('AZURE_BATCHAI_STORAGE_KEY', None)
        with self.assertRaises(CLIError):
            add_azure_file_share_to_cluster_create_parameters(
                params, 'share', 'relative_path')

        os.environ['AZURE_BATCHAI_STORAGE_ACCOUNT'] = 'account'
        os.environ['AZURE_BATCHAI_STORAGE_KEY'] = 'key'

        # No relative mount path provided.
        with self.assertRaises(CLIError):
            add_azure_file_share_to_cluster_create_parameters(
                params, 'share', '')

        # Check valid update.
        result = add_azure_file_share_to_cluster_create_parameters(
            params, 'share', 'relative_path')
        self.assertEquals(
            result.node_setup.mount_volumes.azure_file_shares[0].account_name,
            'account')
        self.assertEquals(
            result.node_setup.mount_volumes.azure_file_shares[0].
            azure_file_url, 'https://account.file.core.windows.net/share')
        self.assertEquals(
            result.node_setup.mount_volumes.azure_file_shares[0].
            relative_mount_path, 'relative_path')
        self.assertEquals(
            result.node_setup.mount_volumes.azure_file_shares[0].credentials.
            account_key, 'key')

        # Account name and key provided via command line args.
        os.environ.pop('AZURE_BATCHAI_STORAGE_ACCOUNT', None)
        os.environ.pop('AZURE_BATCHAI_STORAGE_KEY', None)
        result = add_azure_file_share_to_cluster_create_parameters(
            params, 'share', 'relative_path', 'account', 'key')
        self.assertEquals(
            result.node_setup.mount_volumes.azure_file_shares[0].account_name,
            'account')
        self.assertEquals(
            result.node_setup.mount_volumes.azure_file_shares[0].
            azure_file_url, 'https://account.file.core.windows.net/share')
        self.assertEquals(
            result.node_setup.mount_volumes.azure_file_shares[0].
            relative_mount_path, 'relative_path')
        self.assertEquals(
            result.node_setup.mount_volumes.azure_file_shares[0].credentials.
            account_key, 'key')
 def test_batchai_update_cluster_user_account_settings_using_current_user_and_default_ssh_key(
         self, get_default_ssh_public_key_location, getuser):
     get_default_ssh_public_key_location.return_value = _data_file('key.txt')
     getuser.return_value = 'current_user'
     params = ClusterCreateParameters()
     actual = _update_user_account_settings(params, None, None, None)
     self.assertEquals(actual.user_account_settings,
                       UserAccountSettings(
                           admin_user_name='current_user',
                           admin_user_ssh_public_key=SSH_KEY))
Exemple #10
0
    def test_batchai_add_nfs_to_cluster_create_parameters(self):
        params = ClusterCreateParameters(location='eastus', vm_size='STANDARD_D1',
                                         user_account_settings=UserAccountSettings(admin_user_name='name',
                                                                                   admin_user_password='******'))

        # No relative mount path provided.
        with self.assertRaises(CLIError):
            _add_nfs_to_cluster_create_parameters(params, 'id', '')

        # Check valid update.
        result = _add_nfs_to_cluster_create_parameters(params, 'id', 'relative_path')
        self.assertEquals(result.node_setup.mount_volumes.file_servers[0].file_server.id, 'id')
        self.assertEquals(result.node_setup.mount_volumes.file_servers[0].relative_mount_path, 'relative_path')
        self.assertEquals(result.node_setup.mount_volumes.file_servers[0].mount_options, 'rw')
Exemple #11
0
    def test_batchai_add_azure_container_to_cluster_create_parameters(self):
        cli_ctx = TestCli()
        params = ClusterCreateParameters(location='eastus', vm_size='STANDARD_D1',
                                         user_account_settings=UserAccountSettings(admin_user_name='name',
                                                                                   admin_user_password='******'))

        # No environment variables given.
        os.environ.pop('AZURE_BATCHAI_STORAGE_ACCOUNT', None)
        os.environ.pop('AZURE_BATCHAI_STORAGE_KEY', None)
        with self.assertRaises(CLIError):
            _add_azure_container_to_cluster_create_parameters(cli_ctx, params, 'container', 'relative_path')

        os.environ['AZURE_BATCHAI_STORAGE_ACCOUNT'] = 'account'
        os.environ['AZURE_BATCHAI_STORAGE_KEY'] = 'key'

        # No relative mount path provided.
        with self.assertRaises(CLIError):
            _add_azure_container_to_cluster_create_parameters(cli_ctx, params, 'container', '')

        # Check valid update.
        result = _add_azure_container_to_cluster_create_parameters(cli_ctx, params, 'container', 'relative_path')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].account_name, 'account')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].container_name,
                          'container')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].relative_mount_path,
                          'relative_path')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].credentials.account_key,
                          'key')

        # Account name and key provided via command line args.
        os.environ.pop('AZURE_BATCHAI_STORAGE_ACCOUNT', None)
        os.environ.pop('AZURE_BATCHAI_STORAGE_KEY', None)
        result = _add_azure_container_to_cluster_create_parameters(cli_ctx, params, 'container', 'relative_path',
                                                                   'account', 'key')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].account_name, 'account')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].container_name,
                          'container')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].relative_mount_path,
                          'relative_path')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].credentials.account_key,
                          'key')
Exemple #12
0
    def test_batchai_update_cluster_create_parameters_with_user_account_settings(
            self):
        """Test updating of user account settings."""
        params = ClusterCreateParameters(
            location='eastus',
            vm_size='STANDARD_D1',
            user_account_settings=UserAccountSettings(
                admin_user_name='name', admin_user_password='******'))

        # No update.
        update_user_account_settings(params, None, None, None)
        self.assertEquals(params.user_account_settings.admin_user_name, 'name')
        self.assertEquals(params.user_account_settings.admin_user_password,
                          'password')
        self.assertIsNone(
            params.user_account_settings.admin_user_ssh_public_key)

        # Updating.
        params.user_account_settings = None
        update_user_account_settings(params, 'user', SSH_KEY, None)
        self.assertEquals(params.user_account_settings.admin_user_name, 'user')
        self.assertIsNone(params.user_account_settings.admin_user_password)
        self.assertEquals(
            params.user_account_settings.admin_user_ssh_public_key, SSH_KEY)

        # Incorrect ssh public key.
        params.user_account_settings = None  # user may emit user account settings in config file
        with self.assertRaises(CLIError):
            update_user_account_settings(params, 'user', 'wrong' + SSH_KEY,
                                         'password')

        # No user account.
        params.user_account_settings = None
        with self.assertRaises(CLIError):
            update_user_account_settings(params, None, SSH_KEY, None)

        # No credentials.
        params.user_account_settings = None
        with self.assertRaises(CLIError):
            update_user_account_settings(params, 'user', None, None)

        # ssh public key from a file.
        update_user_account_settings(params, 'user', _data_file('key.txt'),
                                     None)
        self.assertEquals(
            params.user_account_settings.admin_user_ssh_public_key, SSH_KEY)
Exemple #13
0
    def test_batchai_add_azure_container_to_cluster_create_parameters(self):
        """Test adding of azure file share into cluster create parameters."""
        params = ClusterCreateParameters(
            location='eastus',
            vm_size='STANDARD_D1',
            user_account_settings=UserAccountSettings(
                admin_user_name='name', admin_user_password='******'))

        # No environment variables given.
        os.environ.pop('AZURE_BATCHAI_STORAGE_ACCOUNT', None)
        os.environ.pop('AZURE_BATCHAI_STORAGE_KEY', None)
        with self.assertRaises(CLIError):
            add_azure_container_to_cluster_create_parameters(
                params, 'container', 'relative_path')

        os.environ['AZURE_BATCHAI_STORAGE_ACCOUNT'] = 'account'
        os.environ['AZURE_BATCHAI_STORAGE_KEY'] = 'key'

        # No relative mount path provided.
        with self.assertRaises(CLIError):
            add_azure_container_to_cluster_create_parameters(
                params, 'container', '')

        # Check valid update.
        add_azure_container_to_cluster_create_parameters(
            params, 'container', 'relative_path')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_blob_file_systems[0].
            account_name, 'account')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_blob_file_systems[0].
            container_name, 'container')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_blob_file_systems[0].
            relative_mount_path, 'relative_path')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_blob_file_systems[0].
            credentials_info.account_key, 'key')
Exemple #14
0
    def test_batchai_cluster_parameter_update_with_environment_variables(self):
        cli_ctx = TestCli()
        params = ClusterCreateParameters(
            location='eastus', vm_size='STANDARD_D1',
            user_account_settings=UserAccountSettings(admin_user_name='name',
                                                      admin_user_password='******'),
            node_setup=NodeSetup(mount_volumes=MountVolumes(
                azure_file_shares=[AzureFileShareReference(
                    relative_mount_path='azfiles',
                    account_name='<AZURE_BATCHAI_STORAGE_ACCOUNT>',
                    azure_file_url='https://<AZURE_BATCHAI_STORAGE_ACCOUNT>.file.core.windows.net/share',
                    credentials=AzureStorageCredentialsInfo(account_key='<AZURE_BATCHAI_STORAGE_KEY>')
                )],
                azure_blob_file_systems=[AzureBlobFileSystemReference(
                    relative_mount_path='blobfs',
                    container_name='container',
                    account_name='<AZURE_BATCHAI_STORAGE_ACCOUNT>',
                    credentials=AzureStorageCredentialsInfo(account_key='<AZURE_BATCHAI_STORAGE_KEY>')
                )]
            )))

        # No environment variables and no command line parameters provided.
        os.environ.pop('AZURE_BATCHAI_STORAGE_ACCOUNT', None)
        os.environ.pop('AZURE_BATCHAI_STORAGE_KEY', None)
        with self.assertRaises(CLIError):
            _update_cluster_create_parameters_with_env_variables(cli_ctx, params)

        # Set environment variables and check patching results.
        os.environ['AZURE_BATCHAI_STORAGE_ACCOUNT'] = 'account'
        os.environ['AZURE_BATCHAI_STORAGE_KEY'] = 'key'
        result = _update_cluster_create_parameters_with_env_variables(cli_ctx, params)
        self.assertEquals(result.node_setup.mount_volumes.azure_file_shares[0].account_name, 'account')
        self.assertEquals(result.node_setup.mount_volumes.azure_file_shares[0].credentials.account_key, 'key')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].account_name, 'account')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].credentials.account_key, 'key')

        # Test a case when no patching required.
        params.node_setup.mount_volumes.azure_file_shares[0].account_name = 'some_account'
        params.node_setup.mount_volumes.azure_file_shares[0].credentials.account_key = 'some_key'
        params.node_setup.mount_volumes.azure_blob_file_systems[0].account_name = 'some_other_account'
        params.node_setup.mount_volumes.azure_blob_file_systems[0].credentials.account_key = 'some_other_key'
        os.environ['AZURE_BATCHAI_STORAGE_ACCOUNT'] = 'account'
        os.environ['AZURE_BATCHAI_STORAGE_KEY'] = 'key'
        result = _update_cluster_create_parameters_with_env_variables(cli_ctx, params)
        self.assertEquals(result.node_setup.mount_volumes.azure_file_shares[0].account_name, 'some_account')
        self.assertEquals(result.node_setup.mount_volumes.azure_file_shares[0].credentials.account_key, 'some_key')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].account_name,
                          'some_other_account')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].credentials.account_key,
                          'some_other_key')

        # Storage account and key provided as command line args.
        params.node_setup.mount_volumes.azure_file_shares[0].account_name = '<AZURE_BATCHAI_STORAGE_ACCOUNT>'
        params.node_setup.mount_volumes.azure_file_shares[0].credentials.account_key = '<AZURE_BATCHAI_STORAGE_KEY>'
        params.node_setup.mount_volumes.azure_blob_file_systems[0].account_name = '<AZURE_BATCHAI_STORAGE_ACCOUNT>'
        params.node_setup.mount_volumes.azure_blob_file_systems[0].credentials.account_key = \
            '<AZURE_BATCHAI_STORAGE_KEY>'
        result = _update_cluster_create_parameters_with_env_variables(cli_ctx, params, 'account_from_cmd', 'key_from_cmd')
        self.assertEquals(result.node_setup.mount_volumes.azure_file_shares[0].account_name, 'account_from_cmd')
        self.assertEquals(result.node_setup.mount_volumes.azure_file_shares[0].credentials.account_key, 'key_from_cmd')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].account_name, 'account_from_cmd')
        self.assertEquals(result.node_setup.mount_volumes.azure_blob_file_systems[0].credentials.account_key,
                          'key_from_cmd')

        # Non existing storage account provided as command line arg.
        params.node_setup.mount_volumes.azure_file_shares[0].account_name = '<AZURE_BATCHAI_STORAGE_ACCOUNT>'
        params.node_setup.mount_volumes.azure_file_shares[0].credentials.account_key = '<AZURE_BATCHAI_STORAGE_KEY>'
        with self.assertRaises(CLIError):
            _update_cluster_create_parameters_with_env_variables(cli_ctx, params, str(uuid.uuid4()), None)
 def test_batchai_update_cluster_user_account_settings_using_command_line_when_key_is_wrong(self):
     params = ClusterCreateParameters()
     with self.assertRaisesRegexp(CLIError, 'Incorrect ssh public key value'):
         _update_user_account_settings(params, 'user', 'my key', None)
Exemple #16
0
    def test_batchai_cluster_parameter_update_with_environment_variables(self):
        """Test patching of cluster create parameters with environment variables."""
        params = ClusterCreateParameters(
            location='eastus',
            vm_size='STANDARD_D1',
            user_account_settings=UserAccountSettings(
                admin_user_name='name', admin_user_password='******'),
            node_setup=NodeSetup(
                mount_volumes=MountVolumes(
                    azure_file_shares=[
                        AzureFileShareReference(
                            relative_mount_path='azfiles',
                            account_name='<AZURE_BATCHAI_STORAGE_ACCOUNT>',
                            azure_file_url=
                            'https://<AZURE_BATCHAI_STORAGE_ACCOUNT>.file.core.windows.net/share',
                            credentials_info=AzureStorageCredentialsInfo(
                                account_key='<AZURE_BATCHAI_STORAGE_KEY>'))
                    ],
                    azure_blob_file_systems
                    =[
                        AzureBlobFileSystemReference(
                            relative_mount_path='blobfs',
                            container_name='container',
                            account_name='<AZURE_BATCHAI_STORAGE_ACCOUNT>',
                            credentials_info=AzureStorageCredentialsInfo(
                                account_key='<AZURE_BATCHAI_STORAGE_KEY>'))
                    ])))

        # No environment variables provided.
        os.environ.pop('AZURE_BATCHAI_STORAGE_ACCOUNT', None)
        os.environ.pop('AZURE_BATCHAI_STORAGE_KEY', None)
        with self.assertRaises(CLIError):
            update_cluster_create_parameters_with_env_variables(params)

        # Set environment variables and check patching results.
        os.environ['AZURE_BATCHAI_STORAGE_ACCOUNT'] = 'account'
        os.environ['AZURE_BATCHAI_STORAGE_KEY'] = 'key'
        update_cluster_create_parameters_with_env_variables(params)
        self.assertEquals(
            params.node_setup.mount_volumes.azure_file_shares[0].account_name,
            'account')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_file_shares[0].
            credentials_info.account_key, 'key')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_blob_file_systems[0].
            account_name, 'account')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_blob_file_systems[0].
            credentials_info.account_key, 'key')

        # Test a case when no patching required.
        os.environ['AZURE_BATCHAI_STORAGE_ACCOUNT'] = 'another_account'
        os.environ['AZURE_BATCHAI_STORAGE_KEY'] = 'another_key'
        update_cluster_create_parameters_with_env_variables(params)

        # Check that no patching has been done.
        self.assertEquals(
            params.node_setup.mount_volumes.azure_file_shares[0].account_name,
            'account')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_file_shares[0].
            credentials_info.account_key, 'key')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_blob_file_systems[0].
            account_name, 'account')
        self.assertEquals(
            params.node_setup.mount_volumes.azure_blob_file_systems[0].
            credentials_info.account_key, 'key')
    def test_batchai_update_nodes_information(self):
        params = ClusterCreateParameters(location='southcentralus', vm_size='STANDARD_D1',
                                         user_account_settings=UserAccountSettings(admin_user_name='name',
                                                                                   admin_user_password='******'))
        # Update to autoscale Ubuntu DSVM.
        result = _update_nodes_information(params, 'ubuntudsvm', None, 'Standard_NC6', 'dedicated', None, 1, 3)
        self.assertEquals(result.vm_size, 'Standard_NC6')
        self.assertEquals(result.vm_priority, 'dedicated')
        self.assertEquals(
            ImageReference(publisher='microsoft-ads', offer='linux-data-science-vm-ubuntu', sku='linuxdsvmubuntu'),
            result.virtual_machine_configuration.image_reference)
        self.assertEquals(result.scale_settings,
                          ScaleSettings(auto_scale=AutoScaleSettings(minimum_node_count=1, maximum_node_count=3)))

        # Update to manual scale Ubuntu LTS.
        result = _update_nodes_information(params, 'UbuntuLTS', None, 'Standard_NC6', '', None, 2, 2)
        self.assertEquals(result.vm_size, 'Standard_NC6')
        self.assertEquals(result.vm_priority, 'dedicated')
        self.assertEquals(ImageReference(
            publisher='Canonical', offer='UbuntuServer', sku='16.04-LTS'),
            result.virtual_machine_configuration.image_reference
        )
        self.assertEquals(result.scale_settings, ScaleSettings(
            manual=ManualScaleSettings(target_node_count=2)))

        # Update to manual scale Ubuntu LTS.
        result = _update_nodes_information(params, 'UbuntuLTS', None, 'Standard_NC6', '', 2, 2, 2)
        self.assertEquals(result.vm_size, 'Standard_NC6')
        self.assertEquals(result.vm_priority, 'dedicated')
        self.assertEquals(
            ImageReference(publisher='Canonical', offer='UbuntuServer', sku='16.04-LTS'),
            result.virtual_machine_configuration.image_reference)
        self.assertEquals(result.scale_settings, ScaleSettings(manual=ManualScaleSettings(target_node_count=2)))

        # Update to manual scale Ubuntu LTS.
        result = _update_nodes_information(params, 'UbuntuLTS', None, 'Standard_NC6', '', 2, None, None)
        self.assertEquals(result.vm_size, 'Standard_NC6')
        self.assertEquals(result.vm_priority, 'dedicated')
        self.assertEquals(
            ImageReference(publisher='Canonical', offer='UbuntuServer', sku='16.04-LTS'),
            result.virtual_machine_configuration.image_reference
        )
        self.assertEquals(result.scale_settings, ScaleSettings(manual=ManualScaleSettings(target_node_count=2)))

        # Update to auto-scale with initial count Ubuntu LTS.
        result = _update_nodes_information(params, 'UbuntuLTS', None, 'Standard_NC6', '', 1, 0, 2)
        self.assertEquals(result.vm_size, 'Standard_NC6')
        self.assertEquals(result.vm_priority, 'dedicated')
        self.assertEquals(
            ImageReference(publisher='Canonical', offer='UbuntuServer', sku='16.04-LTS'),
            result.virtual_machine_configuration.image_reference)
        self.assertEquals(
            ScaleSettings(
                auto_scale=AutoScaleSettings(minimum_node_count=0, maximum_node_count=2, initial_node_count=1)),
            result.scale_settings)

        # Update image.
        params.scale_settings = ScaleSettings(manual=ManualScaleSettings(target_node_count=2))
        result = _update_nodes_information(params, 'UbuntuDsvm', None, None, '', None, None, None)
        self.assertEquals(
            ImageReference(publisher='microsoft-ads', offer='linux-data-science-vm-ubuntu', sku='linuxdsvmubuntu'),
            result.virtual_machine_configuration.image_reference)
        self.assertEquals(result.scale_settings, ScaleSettings(manual=ManualScaleSettings(target_node_count=2)))

        # Update nothing.
        result = _update_nodes_information(params, None, None, None, '', None, None, None)
        self.assertEqual(params, result)

        # Wrong image.
        with self.assertRaisesRegexp(CLIError, 'Unsupported image alias'):
            _update_nodes_information(params, 'unsupported', None, None, '', None, None, None)

        # No VM size.
        params.vm_size = None
        with self.assertRaisesRegexp(CLIError, 'Please provide VM size'):
            _update_nodes_information(params, 'unsupported', None, None, '', None, None, None)

        # No scale settings.
        params.vm_size = 'Standard_NC6'
        params.scale_settings = None
        with self.assertRaisesRegexp(CLIError, 'Please provide scale setting for the cluster'):
            _update_nodes_information(params, None, None, None, '', None, None, None)

        # Error if only min is specified
        with self.assertRaisesRegexp(CLIError, 'You need to either provide both min and max node'):
            _update_nodes_information(params, 'UbuntuLTS', None, 'Standard_NC6', '', 2, 2, None)
 def test_setup_task_no_output(self):
     with self.assertRaisesRegexp(CLIError, '--setup-task requires providing of --setup-task-output'):
         _add_setup_task('echo hi', None, ClusterCreateParameters())