def test_batchai_add_azure_container_mount_volumes_account_and_key_via_cmd_line_args(self):
actual = _add_azure_container_to_mount_volumes(
TestCli(), MountVolumes(), 'container', 'relative_path', 'account', 'key')
expected = MountVolumes(azure_blob_file_systems=[
AzureBlobFileSystemReference(
account_name='account',
container_name='container',
relative_mount_path='relative_path',
credentials=AzureStorageCredentialsInfo(account_key='key'))])
self.assertEqual(expected, actual)
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_container_mount_volumes_account_and_key_via_env_variables(self):
with _given_env_variable('AZURE_BATCHAI_STORAGE_ACCOUNT', 'account'):
with _given_env_variable('AZURE_BATCHAI_STORAGE_KEY', 'key'):
actual = _add_azure_container_to_mount_volumes(TestCli(), MountVolumes(), 'container', 'relative_path')
expected = MountVolumes(azure_blob_file_systems=[
AzureBlobFileSystemReference(
account_name='account',
container_name='container',
relative_mount_path='relative_path',
credentials=AzureStorageCredentialsInfo(account_key='key'))])
self.assertEqual(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_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_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)