def test_list_setup_task_files_for_cluster_when_output_on_afs(self, get_files_from_afs):
cluster = Cluster(
node_setup=NodeSetup(
setup_task=SetupTask(
comman_line='true',
std_out_err_path_prefix='$AZ_BATCHAI_MOUNT_ROOT/afs')))
cluster.node_setup.setup_task.std_out_err_path_suffix = 'path/segment'
cluster.node_setup.mount_volumes = MountVolumes(
azure_file_shares=[
AzureFileShareReference(relative_mount_path='afs')])
cli_ctx = TestCli()
_list_node_setup_files_for_cluster(cli_ctx, cluster, '.', 60)
get_files_from_afs.assert_called_once_with(
cli_ctx,
AzureFileShareReference(relative_mount_path='afs'),
os.path.join('path/segment', '.'), 60)
def test_batchai_patch_mount_volumes_with_credentials(self, get_storage_client):
# noinspection PyTypeChecker
mount_volumes = MountVolumes(
azure_file_shares=[
AzureFileShareReference(
relative_mount_path='azfiles',
account_name=None,
azure_file_url='https://account1.file.core.windows.net/share',
credentials=None
)
],
azure_blob_file_systems=[
AzureBlobFileSystemReference(
relative_mount_path='blobfs',
container_name='container',
account_name='account2',
credentials=None
),
]
)
get_storage_client.return_value = \
_get_mock_storage_accounts_and_keys({'account1': 'key1', 'account2': 'key2'})
actual = _patch_mount_volumes(TestCli(), mount_volumes)
expected = MountVolumes(
azure_file_shares=[
AzureFileShareReference(
relative_mount_path='azfiles',
account_name='account1',
azure_file_url='https://account1.file.core.windows.net/share',
credentials=AzureStorageCredentialsInfo(account_key='key1')
)
],
azure_blob_file_systems=[
AzureBlobFileSystemReference(
relative_mount_path='blobfs',
container_name='container',
account_name='account2',
credentials=AzureStorageCredentialsInfo(account_key='key2')
)
]
)
self.assertEquals(expected, actual)
def test_batchai_add_azure_file_share_to_absent_mount_volumes(self, get_storage_client):
get_storage_client.return_value = _get_mock_storage_accounts_and_keys({'account': 'key'})
actual = _add_azure_file_share_to_mount_volumes(TestCli(), None, 'share', 'relative_path', 'account')
expected = MountVolumes(azure_file_shares=[
AzureFileShareReference(
account_name='account',
azure_file_url='https://account.file.core.windows.net/share',
relative_mount_path='relative_path',
credentials=AzureStorageCredentialsInfo(
account_key='key'
)
)]
)
self.assertEqual(expected, actual)
def test_batchai_patch_mount_volumes_ill_formed_azure_file_url(self):
# noinspection PyTypeChecker
mount_volumes = MountVolumes(
azure_file_shares=[
AzureFileShareReference(
relative_mount_path='azfiles',
account_name=None,
azure_file_url='http://http://something/is/wrong',
credentials=None
)
]
)
with self.assertRaisesRegexp(CLIError, 'Ill-formed Azure File URL'):
_patch_mount_volumes(TestCli(), mount_volumes)
def test_batchai_patch_mount_volumes_no_azure_file_url(self):
# noinspection PyTypeChecker
mount_volumes = MountVolumes(
azure_file_shares=[
AzureFileShareReference(
relative_mount_path='azfiles',
account_name=None,
azure_file_url=None,
credentials=None
)
]
)
with self.assertRaisesRegexp(CLIError, 'Azure File URL cannot absent or be empty'):
_patch_mount_volumes(TestCli(), mount_volumes)
def test_batchai_patch_mount_volumes_with_credentials_no_account_found(self, get_storage_client):
get_storage_client.return_value = _get_mock_storage_accounts_and_keys({})
# noinspection PyTypeChecker
mount_volumes = MountVolumes(
azure_file_shares=[
AzureFileShareReference(
relative_mount_path='azfiles',
account_name=None,
azure_file_url='https://account1.file.core.windows.net/share',
credentials=None
)
]
)
with self.assertRaisesRegexp(CLIError, 'Cannot find "account1" storage account'):
_patch_mount_volumes(TestCli(), mount_volumes)
def test_list_setup_task_files_for_cluster_when_output_not_on_afs_or_bfs(self):
with self.assertRaisesRegexp(CLIError, 'List files is supported only for clusters with startup task'):
cluster = Cluster(
node_setup=NodeSetup(
setup_task=SetupTask(
comman_line='true',
std_out_err_path_prefix='$AZ_BATCHAI_MOUNT_ROOT/nfs')),
mount_volumes=MountVolumes(
azure_file_shares=[
AzureFileShareReference(
azure_file_url='https://account.file.core.windows.net/share',
relative_mount_path='afs'
)]
))
cluster.node_setup.setup_task.std_out_err_path_suffix = 'path/segment'
_list_node_setup_files_for_cluster(TestCli(), cluster, '.', 60)
def test_batchai_add_azure_file_share_to_mount_volumes_account_and_key_via_env_variables(
self, get_storage_client):
get_storage_client.return_value = _get_mock_storage_accounts_and_keys({'account': 'key'})
with _given_env_variable('AZURE_BATCHAI_STORAGE_ACCOUNT', 'account'):
with _given_env_variable('AZURE_BATCHAI_STORAGE_KEY', 'key'):
actual = _add_azure_file_share_to_mount_volumes(TestCli(), MountVolumes(), 'share', 'relative_path')
expected = MountVolumes(azure_file_shares=[
AzureFileShareReference(
account_name='account',
azure_file_url='https://account.file.core.windows.net/share',
relative_mount_path='relative_path',
credentials=AzureStorageCredentialsInfo(
account_key='key'
)
)]
)
self.assertEqual(expected, actual)
def test_batchai_patch_mount_volumes_when_no_patching_required(
self, get_storage_client):
get_storage_client.return_value = _get_mock_storage_accounts_and_keys(
{})
# noinspection PyTypeChecker
mount_volumes = MountVolumes(
azure_file_shares=[
AzureFileShareReference(
relative_mount_path='azfiles',
account_name='account1',
azure_file_url='https://account1.file.core.windows.net/share',
credentials=AzureStorageCredentialsInfo(account_key='key')
)
]
)
actual = _patch_mount_volumes(TestCli(), mount_volumes)
self.assertEquals(mount_volumes, actual)
def test_batchai_patch_mount_volumes_with_credentials_no_account_given(
self, get_storage_client):
get_storage_client.return_value = _get_mock_storage_accounts_and_keys(
{})
# noinspection PyTypeChecker
mount_volumes = MountVolumes(
azure_file_shares=[
AzureFileShareReference(
relative_mount_path='azfiles',
account_name=None,
azure_file_url='https://<AZURE_BATCHAI_STORAGE_ACCOUNT>.file.core.windows.net/share',
credentials=None
)
]
)
with self.assertRaisesRegexp(CLIError, 'Please configure Azure Storage account name'):
_patch_mount_volumes(TestCli(), mount_volumes)
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_patch_mount_volumes_with_templates_via_env_variables(self):
# noinspection PyTypeChecker
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>')
),
AzureFileShareReference(
relative_mount_path='azfiles2',
account_name=None,
azure_file_url='https://<AZURE_BATCHAI_STORAGE_ACCOUNT>.file.core.windows.net/share2',
credentials=None
)
],
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>')
),
AzureBlobFileSystemReference(
relative_mount_path='blobfs2',
container_name='container2',
account_name='<AZURE_BATCHAI_STORAGE_ACCOUNT>',
credentials=None
),
]
)
with _given_env_variable('AZURE_BATCHAI_STORAGE_ACCOUNT', 'account'):
with _given_env_variable('AZURE_BATCHAI_STORAGE_KEY', 'key'):
actual = _patch_mount_volumes(TestCli(), mount_volumes)
expected = MountVolumes(
azure_file_shares=[
AzureFileShareReference(
relative_mount_path='azfiles',
account_name='account',
azure_file_url='https://account.file.core.windows.net/share',
credentials=AzureStorageCredentialsInfo(account_key='key')
),
AzureFileShareReference(
relative_mount_path='azfiles2',
account_name='account',
azure_file_url='https://account.file.core.windows.net/share2',
credentials=AzureStorageCredentialsInfo(account_key='key')
),
],
azure_blob_file_systems=[
AzureBlobFileSystemReference(
relative_mount_path='blobfs',
container_name='container',
account_name='account',
credentials=AzureStorageCredentialsInfo(account_key='key')
),
AzureBlobFileSystemReference(
relative_mount_path='blobfs2',
container_name='container2',
account_name='account',
credentials=AzureStorageCredentialsInfo(account_key='key')
)
]
)
self.assertEquals(expected, actual)
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)
