def test_batchai_update_cluster_user_account_settings_using_command_line_overriding_config(self): params = ClusterCreateParameters( user_account_settings=UserAccountSettings( admin_user_name='old_name', admin_user_ssh_public_key='old key', admin_user_password='******')) 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_batchai_update_file_server_user_account_settings_using_command_line_when_key_is_path(self): params = FileServerCreateParameters() actual = _update_user_account_settings(params, 'user', _data_file('key.txt'), None) self.assertEquals(actual.ssh_configuration.user_account_settings, UserAccountSettings( admin_user_name='user', admin_user_ssh_public_key=SSH_KEY))
def test_batchai_update_file_server_user_account_settings_using_config(self): params = FileServerCreateParameters( ssh_configuration=SshConfiguration( 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_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_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_file_server_user_account_settings_using_command_line(self): params = FileServerCreateParameters() actual = _update_user_account_settings(params, 'name', SSH_KEY, 'password') self.assertEquals(actual.ssh_configuration.user_account_settings, UserAccountSettings( admin_user_name='name', admin_user_ssh_public_key=SSH_KEY, admin_user_password='******'))
def test_batchai_update_file_server_user_account_settings_using_current_user(self, getuser): getuser.return_value = 'current_user' params = FileServerCreateParameters() actual = _update_user_account_settings(params, None, None, 'password') self.assertEquals(actual.ssh_configuration.user_account_settings, UserAccountSettings( admin_user_name='current_user', admin_user_password='******'))
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_file_server_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 = FileServerCreateParameters() actual = _update_user_account_settings(params, None, None, None) self.assertEquals(actual.ssh_configuration.user_account_settings, UserAccountSettings( admin_user_name='current_user', admin_user_ssh_public_key=SSH_KEY))
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')
def test_batchai_update_file_server_create_parameters_with_user_account_settings(self): params = FileServerCreateParameters(location='eastus', vm_size='STANDARD_D1', ssh_configuration=SshConfiguration( user_account_settings=UserAccountSettings( admin_user_name='name', admin_user_password='******')), data_disks=DataDisks(10, 2, 'Standard_LRS')) # No update. result = _update_user_account_settings(params, None, None, None) self.assertEqual(params, result) # Updating when user_account_setting are omitted. params.ssh_configuration.user_account_settings = None result = _update_user_account_settings(params, 'user', SSH_KEY, None) self.assertEquals(result.ssh_configuration.user_account_settings.admin_user_name, 'user') self.assertIsNone(result.ssh_configuration.user_account_settings.admin_user_password) self.assertEquals(result.ssh_configuration.user_account_settings.admin_user_ssh_public_key, SSH_KEY) # Updating when ssh_configuration is omitted. params.ssh_configuration = None result = _update_user_account_settings(params, 'user', SSH_KEY, 'password') self.assertEquals(result.ssh_configuration.user_account_settings.admin_user_name, 'user') self.assertEquals(result.ssh_configuration.user_account_settings.admin_user_password, 'password') self.assertEquals(result.ssh_configuration.user_account_settings.admin_user_ssh_public_key, SSH_KEY) # Incorrect ssh public key. params.ssh_configuration = SshConfiguration(None) # user may emit user account settings in config file with self.assertRaises(CLIError): _update_user_account_settings(params, 'user', 'wrong' + SSH_KEY, None) # No user account. params.ssh_configuration.user_account_settings = None with self.assertRaises(CLIError): _update_user_account_settings(params, None, SSH_KEY, None) # No credentials. params.ssh_configuration.user_account_settings = None with self.assertRaises(CLIError): _update_user_account_settings(params, 'user', None, None) # Only password. params.ssh_configuration.user_account_settings = None result = _update_user_account_settings(params, 'user', None, 'password') self.assertEquals(result.ssh_configuration.user_account_settings.admin_user_name, 'user') self.assertEquals(result.ssh_configuration.user_account_settings.admin_user_password, 'password') # ssh public key from a file. result = _update_user_account_settings(params, 'user', _data_file('key.txt'), None) self.assertEquals(result.ssh_configuration.user_account_settings.admin_user_ssh_public_key, SSH_KEY)
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)
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')
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')
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_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)