# coding: utf-8 #------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. #-------------------------------------------------------------------------- import datetime import io import json import logging import os import sys import time import unittest import requests from testutils.common_recordingtestcase import ( RecordingTestCase, TestMode, ) from tests.mgmt_testcase import HttpStatusCode, AzureMgmtTestCase import tests.mgmt_settings_fake as fake_settings from testutils.common_recordingtestcase import record import azure.mgmt.resource import azure.mgmt.storage import azure.mgmt.batch import azure.batch as batch from azure.batch.batch_auth import SharedKeyCredentials AZURE_BATCH_ACCOUNT = 'pythonsdktest' AZURE_RESOURCE_GROUP = 'python_batch_sdk_test' AZURE_LOCATION = 'westus' AZURE_STORAGE_ACCOUNT = 'pythonsdkbatchtest' LOG = logging.getLogger('batch-python-tests') LOG.level = logging.WARNING LOG.addHandler(logging.StreamHandler()) def init_tst_mode(working_folder): try: path = os.path.join(working_folder, 'testsettings_local.json') with open(path) as testsettings_local_file: test_settings = json.load(testsettings_local_file) return test_settings['mode'] except: return TestMode.playback def create_mgmt_client(settings, client_class, **kwargs): client = client_class( credentials=settings.get_credentials(), subscription_id=settings.SUBSCRIPTION_ID, **kwargs ) return client def create_resource_group(client): group = azure.mgmt.resource.resources.models.ResourceGroup( name=AZURE_RESOURCE_GROUP, location=AZURE_LOCATION ) result_create = client.resource_groups.create_or_update( AZURE_RESOURCE_GROUP, group, ) return result_create def create_storage_account(client): params = azure.mgmt.storage.models.StorageAccountCreateParameters( sku=azure.mgmt.storage.models.Sku( azure.mgmt.storage.models.SkuName.standard_lrs ), kind=azure.mgmt.storage.models.Kind.storage, location=AZURE_LOCATION ) result_create = client.storage_accounts.create( AZURE_RESOURCE_GROUP, AZURE_STORAGE_ACCOUNT, params, ) return result_create.result() def create_batch_account(client, settings, live): if live: storage_resource = '/subscriptions/{}/resourceGroups/{}/providers/Microsoft.Storage/storageAccounts/{}'.format( settings.SUBSCRIPTION_ID, AZURE_RESOURCE_GROUP, AZURE_STORAGE_ACCOUNT ) batch_account = azure.mgmt.batch.models.BatchAccountCreateParameters( location=AZURE_LOCATION, auto_storage=azure.mgmt.batch.models.AutoStorageBaseProperties(storage_resource) ) account_setup = client.account.create(AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, batch_account) new_account = account_setup.result() account_keys = client.account.list_keys(AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT) batch_creds = SharedKeyCredentials(AZURE_BATCH_ACCOUNT, account_keys.primary) else: batch_creds = SharedKeyCredentials(AZURE_BATCH_ACCOUNT, 'ZmFrZV9hY29jdW50X2tleQ==') url = "https://{}.{}.batch.azure.com/".format(AZURE_BATCH_ACCOUNT, AZURE_LOCATION) return azure.batch.BatchServiceClient(batch_creds, base_url=url) class BatchMgmtTestCase(RecordingTestCase): @classmethod def setUpClass(cls): LOG.warning('Starting Batch tests') LOG.debug("Setting up Batch tests:") cls.working_folder = os.path.dirname(__file__) try: cls.test_mode = init_tst_mode(cls.working_folder) cls.fake_settings = fake_settings if TestMode.is_playback(cls.test_mode): LOG.debug(" running in playback mode") cls.live = False cls.settings = cls.fake_settings else: LOG.debug(" running in live mode") import tests.mgmt_settings_real as real_settings cls.settings = real_settings cls.live = True LOG.debug(' creating resource client') cls.resource_client = create_mgmt_client(cls.settings, azure.mgmt.resource.resources.ResourceManagementClient ) LOG.debug(' creating storage client') cls.storage_client = create_mgmt_client(cls.settings, azure.mgmt.storage.StorageManagementClient ) LOG.debug(' creating batch client') cls.batch_mgmt_client = create_mgmt_client(cls.settings, azure.mgmt.batch.BatchManagementClient ) if cls.live: LOG.debug(' creating resource group') create_resource_group(cls.resource_client) LOG.debug(' creating storage') create_storage_account(cls.storage_client) LOG.debug(' creating batch account') cls.batch_client = create_batch_account(cls.batch_mgmt_client, cls.settings, cls.live) except Exception as err: cls.tearDownClass() raise AssertionError("Failed to setup Batch Account: {}".format(err)) LOG.debug(" finished setup") return super(BatchMgmtTestCase, cls).setUpClass() @classmethod def tearDownClass(cls): LOG.debug("Tearing down Batch resources:") if cls.live: try: LOG.debug(" deleting Batch account") deleting = cls.batch_mgmt_client.account.delete( AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT) deleting.wait() except: pass # This should get deleted with the resource group anyway try: LOG.debug(" deleting storage") cls.storage_client.storage_accounts.delete( AZURE_RESOURCE_GROUP, AZURE_STORAGE_ACCOUNT) except: pass try: LOG.debug(" deleting resource group") deleting = cls.resource_client.resource_groups.delete(AZURE_RESOURCE_GROUP) deleting.wait() except: pass LOG.debug(" finished") LOG.warning("Batch tests complete") return super(BatchMgmtTestCase, cls).tearDownClass() def _scrub(self, val): if not hasattr(val, 'replace'): return val val = super(BatchMgmtTestCase, self)._scrub(val) real_to_fake_dict = { self.settings.SUBSCRIPTION_ID: self.fake_settings.SUBSCRIPTION_ID, } val = self._scrub_using_dict(val, real_to_fake_dict) return val def assertBatchError(self, error, message, code, func, *args, **kwargs): try: func(*args, **kwargs) if message not in error: error[message] = "BatchErrorException expected but not raised" except batch.models.BatchErrorException as err: self.assertEqual(error, message, err.error.code, code) except Exception as err: if message not in error: error[message] = "Expected BatchErrorExcption, instead got: {!r}".format(err) def assertRuns(self, error, message, func, *args, **kwargs): try: return func(*args, **kwargs) except Exception as err: if message not in error: error[message] = str(err) def assertList(self, error, message, func, *args, **kwargs): response = self.assertRuns(error, message, func, *args, **kwargs) if response is None: if message not in error: error[message] = "Received uniterable response" return [] try: return [r for r in response] except Exception as err: if message not in error: error[message] = str(err) return [] def assertIsNone(self, error, message, obj): try: super(BatchMgmtTestCase, self).assertIsNone(obj) except AssertionError as err: if message not in error: error[message] = str(err) def assertEqual(self, error, message, first, second, msg = None): try: super(BatchMgmtTestCase, self).assertEqual(first, second) except AssertionError as err: if message not in error: error[message] = str(err) def assertTrue(self, error, message, expr, msg=None): try: super(BatchMgmtTestCase, self).assertTrue(expr, msg) except AssertionError as err: if message not in error: error[message] = str(err) #def assertIsInstance(self, error, message, obj, obj_type): # try: # super(BatchMgmtTestCase, self).assertIsInstance(obj, obj_type) # except AssertionError as err: # if message not in error: # error[message] def assertSuccess(self, errors): if errors: message = "The following errors occurred:\n" message += "\n".join(["{}: {!r}".format(k, v) for k, v in errors.items()]) raise AssertionError(message) @record def test_batch_accounts(self): _e = {} _m = "Test Get Subscription Quota" LOG.debug(_m) quotas = self.assertRuns(_e, _m, self.batch_mgmt_client.subscription.get_subscription_quotas, AZURE_LOCATION) self.assertIsInstance(quotas, azure.mgmt.batch.models.SubscriptionQuotasGetResult) if quotas: self.assertEqual(_e, _m, quotas.account_quota, 1) _m = "Test Get Account" LOG.debug(_m) account = self.assertRuns(_e, _m, self.batch_mgmt_client.account.get, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT) self.assertIsInstance(account, azure.mgmt.batch.models.AccountResource) if account: self.assertEqual(_e, _m, account.core_quota, 20) self.assertEqual(_e, _m, account.pool_quota, 20) _m = "Test List Accounts" #TODO: Need to re-record LOG.debug('TODO: ' + _m) #accounts = self.assertList(_e, _m, self.batch_mgmt_client.account.list) #self.assertTrue(_e, _m, len(accounts) > 0) _m = "Test List Accounts by Resource Group" LOG.debug(_m) accounts = self.assertList(_e, _m, self.batch_mgmt_client.account.list_by_resource_group, AZURE_RESOURCE_GROUP) self.assertEqual(_e, _m, len(accounts), 1) _m = "Test Regenerate Account Key" LOG.debug(_m) keys = self.assertRuns(_e, _m, self.batch_mgmt_client.account.regenerate_key, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'Secondary') self.assertIsInstance(keys, azure.mgmt.batch.models.BatchAccountRegenerateKeyResult) _m = "Test Sync AutoStorage Keys" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_mgmt_client.account.synchronize_auto_storage_keys, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT) self.assertIsNone(_e, _m, response) _m = "Test Update Account" LOG.debug(_m) tags = {'Name': 'tagName', 'Value': 'tagValue'} updated = self.assertRuns(_e, _m, self.batch_mgmt_client.account.update, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, tags) self.assertIsInstance(updated, azure.mgmt.batch.models.AccountResource) if account: self.assertEqual(_e, _m, updated.tags['Name'], 'tagName') self.assertEqual(_e, _m, updated.tags['Value'], 'tagValue') self.assertSuccess(_e) @record def test_batch_applications(self): _e = {} _m = "Test Add Application" LOG.debug(_m) application = self.assertRuns(_e, _m, self.batch_mgmt_client.application.add_application, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id', allow_updated=True, display_name='my_application_name') self.assertIsInstance(application, azure.mgmt.batch.models.Application) if application: self.assertEqual(_e, _m, application.id, 'my_application_id') self.assertEqual(_e, _m, application.display_name, 'my_application_name') self.assertEqual(_e, _m, application.allow_updates, True) _m = "Test Mgmt Get Application" LOG.debug(_m) application = self.assertRuns(_e, _m, self.batch_mgmt_client.application.get_application, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id') self.assertIsInstance(application, azure.mgmt.batch.models.Application) if application: self.assertEqual(_e, _m, application.id, 'my_application_id') self.assertEqual(_e, _m, application.display_name, 'my_application_name') self.assertEqual(_e, _m, application.allow_updates, True) _m = "Test Mgmt List Applications" LOG.debug(_m) applications = self.assertList(_e, _m, self.batch_mgmt_client.application.list, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT) self.assertTrue(_e, _m, len(applications) > 0) _m = "Test Add Application Package" LOG.debug(_m) package_ref = self.assertRuns(_e, _m, self.batch_mgmt_client.application.add_application_package, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id', 'v1.0') self.assertIsInstance(package_ref, azure.mgmt.batch.models.AddApplicationPackageResult) if package_ref: try: with io.BytesIO(b'Hello World') as f: headers = {'x-ms-blob-type': 'BlockBlob'} upload = requests.put(package_ref.storage_url, headers=headers, data=f.read()) if not upload: raise ValueError('Upload failed: {!r}'.format(upload)) except Exception as err: _e[_m] = 'Failed to upload test package: {}'.format(err) _m = "Test Activate Application Package" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_mgmt_client.application.activate_application_package, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id', 'v1.0', 'zip') self.assertIsNone(_e, _m, response) _m = "Test Update Application" LOG.debug(_m) params = azure.mgmt.batch.models.UpdateApplicationParameters( allow_updates=False, display_name='my_updated_name', default_version='v1.0' ) response = self.assertRuns(_e, _m, self.batch_mgmt_client.application.update_application, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id', params) self.assertIsNone(_e, _m, response) _m = "Test Get Application Package" LOG.debug(_m) package_ref = self.assertRuns(_e, _m, self.batch_mgmt_client.application.get_application_package, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id', 'v1.0') self.assertIsInstance(package_ref, azure.mgmt.batch.models.GetApplicationPackageResult) if package_ref: self.assertEqual(_e, _m, package_ref.id, 'my_application_id') self.assertEqual(_e, _m, package_ref.version, 'v1.0') self.assertEqual(_e, _m, package_ref.format, 'zip') self.assertEqual(_e, _m, package_ref.state, azure.mgmt.batch.models.PackageState.active) _m = "Test Service Get Application" LOG.debug(_m) application = self.assertRuns(_e, _m, self.batch_client.application.get, 'my_application_id') self.assertIsInstance(application, batch.models.ApplicationSummary) if application: self.assertEqual(_e, _m, application.id, 'my_application_id') self.assertEqual(_e, _m, application.display_name, 'my_updated_name') self.assertEqual(_e, _m, application.versions, ['v1.0']) _m = "Test Service List Applications" LOG.debug(_m) applications = self.assertList(_e, _m, self.batch_client.application.list) self.assertTrue(_e, _m, len(applications) > 0) _m = "Test Delete Application Package" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_mgmt_client.application.delete_application_package, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id', 'v1.0') self.assertIsNone(_e, _m, response) _m = "Test Delete Application" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_mgmt_client.application.delete_application, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id') self.assertIsNone(_e, _m, response) self.assertSuccess(_e) @record def test_batch_certificates(self): _e = {} _m = "Test Add Certificate" LOG.debug(_m) certificate = batch.models.CertificateAddParameter( thumbprint='cff2ab63c8c955aaf71989efa641b906558d9fb7', thumbprint_algorithm='sha1', data='MIIGMQIBAzCCBe0GCSqGSIb3DQEHAaCCBd4EggXaMIIF1jCCA8AGCSqGSIb3DQEHAaCCA7EEggOtMIIDqTCCA6UGCyqGSIb3DQEMCgECoIICtjCCArIwHAYKKoZIhvcNAQwBAzAOBAhyd3xCtln3iQICB9AEggKQhe5P10V9iV1BsDlwWT561Yu2hVq3JT8ae/ebx1ZR/gMApVereDKkS9Zg4vFyssusHebbK5pDpU8vfAqle0TM4m7wGsRj453ZorSPUfMpHvQnAOn+2pEpWdMThU7xvZ6DVpwhDOQk9166z+KnKdHGuJKh4haMT7Rw/6xZ1rsBt2423cwTrQVMQyACrEkianpuujubKltN99qRoFAxhQcnYE2KlYKw7lRcExq6mDSYAyk5xJZ1ZFdLj6MAryZroQit/0g5eyhoNEKwWbi8px5j71pRTf7yjN+deMGQKwbGl+3OgaL1UZ5fCjypbVL60kpIBxLZwIJ7p3jJ+q9pbq9zSdzshPYor5lxyUfXqaso/0/91ayNoBzg4hQGh618PhFI6RMGjwkzhB9xk74iweJ9HQyIHf8yx2RCSI22JuCMitPMWSGvOszhbNx3AEDLuiiAOHg391mprEtKZguOIr9LrJwem/YmcHbwyz5YAbZmiseKPkllfC7dafFfCFEkj6R2oegIsZo0pEKYisAXBqT0g+6/jGwuhlZcBo0f7UIZm88iA3MrJCjlXEgV5OcQdoWj+hq0lKEdnhtCKr03AIfukN6+4vjjarZeW1bs0swq0l3XFf5RHa11otshMS4mpewshB9iO9MuKWpRxuxeng4PlKZ/zuBqmPeUrjJ9454oK35Pq+dghfemt7AUpBH/KycDNIZgfdEWUZrRKBGnc519C+RTqxyt5hWL18nJk4LvSd3QKlJ1iyJxClhhb/NWEzPqNdyA5cxen+2T9bd/EqJ2KzRv5/BPVwTQkHH9W/TZElFyvFfOFIW2+03RKbVGw72Mr/0xKZ+awAnEfoU+SL/2Gj2m6PHkqFX2sOCi/tN9EA4xgdswEwYJKoZIhvcNAQkVMQYEBAEAAAAwXQYJKwYBBAGCNxEBMVAeTgBNAGkAYwByAG8AcwBvAGYAdAAgAFMAdAByAG8AbgBnACAAQwByAHkAcAB0AG8AZwByAGEAcABoAGkAYwAgAFAAcgBvAHYAaQBkAGUAcjBlBgkqhkiG9w0BCRQxWB5WAFAAdgBrAFQAbQBwADoANABjAGUANgAwADQAZABhAC0AMAA2ADgAMQAtADQANAAxADUALQBhADIAYwBhAC0ANQA3ADcAMwAwADgAZQA2AGQAOQBhAGMwggIOBgkqhkiG9w0BBwGgggH/BIIB+zCCAfcwggHzBgsqhkiG9w0BDAoBA6CCAcswggHHBgoqhkiG9w0BCRYBoIIBtwSCAbMwggGvMIIBXaADAgECAhAdka3aTQsIsUphgIXGUmeRMAkGBSsOAwIdBQAwFjEUMBIGA1UEAxMLUm9vdCBBZ2VuY3kwHhcNMTYwMTAxMDcwMDAwWhcNMTgwMTAxMDcwMDAwWjASMRAwDgYDVQQDEwdub2Rlc2RrMIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQC5fhcxbJHxxBEIDzVOMc56s04U6k4GPY7yMR1m+rBGVRiAyV4RjY6U936dqXHCVD36ps2Q0Z+OeEgyCInkIyVeB1EwXcToOcyeS2YcUb0vRWZDouC3tuFdHwiK1Ed5iW/LksmXDotyV7kpqzaPhOFiMtBuMEwNJcPge9k17hRgRQIDAQABo0swSTBHBgNVHQEEQDA+gBAS5AktBh0dTwCNYSHcFmRjoRgwFjEUMBIGA1UEAxMLUm9vdCBBZ2VuY3mCEAY3bACqAGSKEc+41KpcNfQwCQYFKw4DAh0FAANBAHl2M97QbpzdnwO5HoRBsiEExOcLTNg+GKCr7HUsbzfvrUivw+JLL7qjHAIc5phnK+F5bQ8HKe0L9YXBSKl+fvwxFTATBgkqhkiG9w0BCRUxBgQEAQAAADA7MB8wBwYFKw4DAhoEFGVtyGMqiBd32fGpzlGZQoRM6UQwBBTI0YHFFqTS4Go8CoLgswn29EiuUQICB9A=', certificate_format=batch.models.CertificateFormat.pfx, password='nodesdk') response = self.assertRuns(_e, _m, self.batch_client.certificate.add, certificate) self.assertIsNone(_e, _m, response) _m = "Test List Certificates" LOG.debug(_m) certs = self.assertList(_e, _m, self.batch_client.certificate.list) self.assertTrue(_e, _m, len(certs) > 0) test_cert = [c for c in certs if c.thumbprint == 'cff2ab63c8c955aaf71989efa641b906558d9fb7'] self.assertEqual(_e, _m, len(test_cert), 1) _m = "Test Get Certificate" LOG.debug(_m) cert = self.assertRuns(_e, _m, self.batch_client.certificate.get, 'sha1', 'cff2ab63c8c955aaf71989efa641b906558d9fb7') self.assertIsInstance(cert, batch.models.Certificate) if cert: self.assertEqual(_e, _m, cert.thumbprint, 'cff2ab63c8c955aaf71989efa641b906558d9fb7') self.assertEqual(_e, _m, cert.thumbprint_algorithm, 'sha1') self.assertIsNone(_e, _m, cert.delete_certificate_error) _m = "Test Cancel Certificate Delete" LOG.debug(_m) self.assertBatchError(_e, _m, 'CertificateStateActive', self.batch_client.certificate.cancel_deletion, 'sha1', 'cff2ab63c8c955aaf71989efa641b906558d9fb7') _m = "Test Delete Certificate" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.certificate.delete, 'sha1', 'cff2ab63c8c955aaf71989efa641b906558d9fb7') self.assertIsNone(_e, _m, response) self.assertSuccess(_e) @record def test_batch_pools(self): _e = {} with BatchPool(self.live, self.batch_client, 'python_test_pool_1') as pool_id: with BatchPool(self.live, self.batch_client, 'python_test_pool_2') as pool_id_2: _m = "Test Create Pool with Network Configuration" LOG.debug(_m) pool_config = batch.models.CloudServiceConfiguration('4') network_config = batch.models.NetworkConfiguration('/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/test/providers/Microsoft.Network/virtualNetworks/vnet1/subnets/subnet1') pool = batch.models.PoolAddParameter('no_pool', 'small', cloud_service_configuration=pool_config, network_configuration=network_config) self.assertBatchError(_e, _m, 'InvalidPropertyValue', self.batch_client.pool.add, pool, batch.models.PoolAddOptions(timeout=45)) _m = "Test Upgrade Pool OS" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.pool.upgrade_os, pool_id, 'WA-GUEST-OS-4.32_201605-01') self.assertIsNone(_e, _m, response) _m = "Test Update Pool Parameters" LOG.debug(_m) params = batch.models.PoolUpdatePropertiesParameter([], [], [batch.models.MetadataItem('foo', 'bar')]) response = self.assertRuns(_e, _m, self.batch_client.pool.update_properties, pool_id, params) self.assertIsNone(_e, _m, response) _m = "Test Patch Pool Parameters" LOG.debug(_m) params = batch.models.PoolPatchParameter(metadata=[batch.models.MetadataItem('foo2', 'bar2')]) response = self.assertRuns(_e, _m, self.batch_client.pool.patch, pool_id, params) self.assertIsNone(_e, _m, response) _m = "Test Get Pool" LOG.debug(_m) pool = self.assertRuns(_e, _m, self.batch_client.pool.get, pool_id) self.assertIsInstance(pool, batch.models.CloudPool) if pool: self.assertEqual(_e, _m, pool.id, pool_id) self.assertEqual(_e, _m, pool.state, batch.models.PoolState.upgrading) self.assertEqual(_e, _m, pool.allocation_state, batch.models.AllocationState.steady) self.assertEqual(_e, _m, pool.cloud_service_configuration.os_family, '4') self.assertEqual(_e, _m, pool.vm_size, 'small') self.assertEqual(_e, _m, pool.metadata[0].name, 'foo2') self.assertEqual(_e, _m, pool.metadata[0].value, 'bar2') _m = "Test Get Pool with OData Clauses" LOG.debug(_m) options = batch.models.PoolGetOptions(select='id,state', expand='stats') pool = self.assertRuns(_e, _m, self.batch_client.pool.get, pool_id, options) self.assertIsInstance(pool, batch.models.CloudPool) if pool: self.assertEqual(_e, _m, pool.id, pool_id) self.assertEqual(_e, _m, pool.state, batch.models.PoolState.upgrading) self.assertIsNone(_e, _m, pool.allocation_state) self.assertIsNone(_e, _m, pool.vm_size) _m = "Test Enable Autoscale" LOG.debug(_m) interval = datetime.timedelta(minutes=6) response = self.assertRuns(_e, _m, self.batch_client.pool.enable_auto_scale, pool_id, auto_scale_formula='$TargetDedicated=2', auto_scale_evaluation_interval=interval) self.assertIsNone(_e, _m, response) _m = "Test Evaluate Autoscale" LOG.debug(_m) result = self.assertRuns(_e, _m, self.batch_client.pool.evaluate_auto_scale, pool_id, '$TargetDedicated=3') self.assertIsInstance(result, batch.models.AutoScaleRun) self.assertEqual(_e, _m, result.results, '$TargetDedicated=3;$NodeDeallocationOption=requeue') _m = "Test Disable Autoscale" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.pool.disable_auto_scale, pool_id) self.assertIsNone(_e, _m, response) _m = "Test List Pools without Filters" LOG.debug(_m) pools = self.assertList(_e, _m, self.batch_client.pool.list) self.assertTrue(_e, _m, len(pools) > 1) _m = "Test List Pools with Maximum" LOG.debug(_m) options = batch.models.PoolListOptions(max_results=1) pools = self.assertRuns(_e, _m, self.batch_client.pool.list, options) if pools: self.assertRuns(_e, _m, pools.next) self.assertEqual(_e, _m, len(pools.current_page), 1) _m = "Test List Pools with Filter" LOG.debug(_m) options = batch.models.PoolListOptions( filter='startswith(id,\'python_test_pool_1\')', select='id,state', expand='stats') pools = self.assertList(_e, _m, self.batch_client.pool.list, options) self.assertEqual(_e, _m, len(pools), 1) _m = "Test Pool Exists" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.pool.exists, pool_id_2) self.assertTrue(_e, _m, response) _m = "Test Pool Resize" LOG.debug(_m) params = batch.models.PoolResizeParameter(target_dedicated=3) response = self.assertRuns(_e, _m, self.batch_client.pool.resize, pool_id_2, params) self.assertIsNone(_e, _m, response) _m = "Test Stop Pool Resize" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.pool.stop_resize, pool_id_2) self.assertIsNone(_e, _m, response) _m = "Test Get All Pools Lifetime Statistics" LOG.debug(_m) stats = self.assertRuns(_e, _m, self.batch_client.pool.get_all_pools_lifetime_statistics) self.assertIsInstance(stats, batch.models.PoolStatistics) if stats: self.assertEqual(_e, _m, stats.url, "https://{}.{}.batch.azure.com/lifetimepoolstats".format( AZURE_BATCH_ACCOUNT, AZURE_LOCATION)) self.assertEqual(_e, _m, stats.resource_stats.avg_cpu_percentage, 0.0) self.assertEqual(_e, _m, stats.resource_stats.network_read_gi_b, 0.0) self.assertEqual(_e, _m, stats.resource_stats.disk_write_gi_b, 0.0) self.assertEqual(_e, _m, stats.resource_stats.peak_disk_gi_b, 0.0) _m = "Test Get Pool Usage Info" #TODO: Test with real usage metrics LOG.debug('TODO: ' + _m) response = self.assertList(_e, _m, self.batch_client.pool.list_pool_usage_metrics) #TODO: Assert self.assertSuccess(_e) @record def test_batch_compute_nodes(self): _e = {} with BatchPool(self.live, self.batch_client, 'python_test_pool_3', target_dedicated=2) as pool_id: _m = "Test List Compute Nodes" LOG.debug(_m) nodes = self.assertList(_e, _m, self.batch_client.compute_node.list, pool_id) self.assertEqual(_e, _m, len(nodes), 2) node_ids = [n.id for n in nodes] _m = "Test Get Compute Node" LOG.debug(_m) node = self.assertRuns(_e, _m, self.batch_client.compute_node.get, pool_id, node_ids[0]) self.assertIsInstance(node, batch.models.ComputeNode) if node: self.assertEqual(_e, _m, node.state, batch.models.ComputeNodeState.idle) self.assertEqual(_e, _m, node.scheduling_state, batch.models.SchedulingState.enabled) _m = "Test Add User" LOG.debug(_m) user = batch.models.ComputeNodeUser('BatchPythonSDKUser', password='kt#_gahr!@aGERDXA', is_admin=False) response = self.assertRuns(_e, _m, self.batch_client.compute_node.add_user, pool_id, node_ids[0], user) self.assertIsNone(_e, _m, response) _m = "Test Update User" LOG.debug(_m) user = batch.models.NodeUpdateUserParameter(password='liilef#$DdRGSa_ewkjh') response = self.assertRuns(_e, _m, self.batch_client.compute_node.update_user, pool_id, node_ids[0], 'BatchPythonSDKUser', user) self.assertIsNone(_e, _m, response) _m = "Test Get RDP File" LOG.debug(_m) file_length = 0 with io.BytesIO() as file_handle: response = self.assertRuns(_e, _m, self.batch_client.compute_node.get_remote_desktop, pool_id, node_ids[0]) if response: for data in response: file_length += len(data) self.assertTrue(_e, _m, file_length > 0) _m = "Test Delete User" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.compute_node.delete_user, pool_id, node_ids[0], 'BatchPythonSDKUser') self.assertIsNone(_e, _m, response) _m = "Test Disable Scheduling" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.compute_node.disable_scheduling, pool_id, node_ids[0]) self.assertIsNone(_e, _m, response) _m = "Test Enable Scehduling" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.compute_node.enable_scheduling, pool_id, node_ids[0]) self.assertIsNone(_e, _m, response) _m = "Test List Files" LOG.debug(_m) files = self.assertList(_e, _m, self.batch_client.file.list_from_compute_node, pool_id, node_ids[0]) self.assertTrue(_e, _m, len(files) > 2) _m = "Test File Properties" LOG.debug('TODO: ' + _m) #props = self.batch_client.file.get_node_file_properties_from_compute_node(pool_id, node_ids[0], '', raw=True) #self.assertTrue('Content-Length' in props.headers) #self.assertTrue('Content-Type'in props.headers) _m = "Test Get File" LOG.debug('TODO: ' + _m) #file_length = 0 #with io.BytesIO() as file_handle: # response = self.batch_client.file.get_from_task(pool_id, node_ids[0], '') # for data in response: # file_length += len(data) #self.assertEqual(file_length, props.headers['Content-Length']) _m = "Test Reboot Node" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.compute_node.reboot, pool_id, node_ids[0]) self.assertIsNone(_e, _m, response) _m = "Test Reimage Node" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.compute_node.reimage, pool_id, node_ids[1]) self.assertIsNone(_e, _m, response) _m = "Test Remove Nodes" LOG.debug(_m) options = batch.models.NodeRemoveParameter(node_ids) response = self.assertRuns(_e, _m, self.batch_client.pool.remove_nodes, pool_id, options) self.assertIsNone(_e, _m, response) self.assertSuccess(_e) @record def test_batch_jobs(self): _e = {} with BatchPool(self.live, self.batch_client, 'python_test_pool_4', target_dedicated=2) as pool_id: _m = "Test Create Job" LOG.debug(_m) job = batch.models.JobAddParameter('python_test_job', batch.models.PoolInformation(pool_id=pool_id)) response = self.assertRuns(_e, _m, self.batch_client.job.add, job) self.assertIsNone(_e, _m, response) LOG.debug(" wait for job to create...") self.sleep(10) _m = "Test Create Job with Auto Complete" LOG.debug(_m) job_with_auto_complete = batch.models.JobAddParameter(id='python_test_job_2', pool_info=batch.models.PoolInformation(pool_id=pool_id), on_all_tasks_complete='noAction', on_task_failure='performExitOptionsJobAction') response = self.assertRuns(_e, _m, self.batch_client.job.add, job_with_auto_complete) self.assertIsNone(_e, _m, response) LOG.debug(" wait for job to create...") self.sleep(10) _m = "Test Update Job" LOG.debug(_m) constraints = batch.models.JobConstraints(max_task_retry_count=3) options = batch.models.JobUpdateParameter(priority=500, constraints=constraints, pool_info=batch.models.PoolInformation(pool_id=pool_id)) response = self.assertRuns(_e, _m, self.batch_client.job.update, 'python_test_job', options) self.assertIsNone(_e, _m, response) _m = "Test Patch Job" LOG.debug(_m) constraints = batch.models.JobConstraints(max_task_retry_count=1) options = batch.models.JobPatchParameter(priority=900, constraints=constraints, pool_info=batch.models.PoolInformation(pool_id=pool_id)) response = self.assertRuns(_e, _m, self.batch_client.job.patch, 'python_test_job', options) self.assertIsNone(_e, _m, response) _m = "Test Get Job" LOG.debug(_m) job = self.assertRuns(_e, _m, self.batch_client.job.get, 'python_test_job') self.assertIsInstance(job, batch.models.CloudJob) if job: self.assertEqual(_e, _m, job.id, 'python_test_job') self.assertEqual(_e, _m, job.pool_info.pool_id, pool_id) self.assertEqual(_e, _m, job.constraints.max_task_retry_count, 1) _m = "Test List Jobs" LOG.debug(_m) jobs = self.assertList(_e, _m, self.batch_client.job.list) self.assertTrue(_e, _m, len(jobs) > 0) job_ids = [j.id for j in jobs] _m = "Test Job Disable" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.job.disable, 'python_test_job', 'requeue') self.assertIsNone(_e, _m, response) _m = "Test Job Enable" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.job.enable, 'python_test_job') self.assertIsNone(_e, _m, response) _m = "Test Create Task with Auto Complete" LOG.debug(_m) conditions = batch.models.ExitConditions(default=batch.models.ExitOptions('terminate'), exit_codes=[batch.models.ExitCodeMapping(code=1, exit_options=batch.models.ExitOptions('None'))]) task = batch.models.TaskAddParameter(id='python_task_with_auto_complete', command_line='cmd /c echo hello world', exit_conditions=conditions) response = self.assertRuns(_e, _m, self.batch_client.task.add, 'python_test_job_2', task) self.assertIsNone(_e, _m, response) _m = "Test Get Task with Auto Complete" LOG.debug(_m) task = self.assertRuns(_e, _m, self.batch_client.task.get, 'python_test_job_2', 'python_task_with_auto_complete') self.assertIsInstance(task, batch.models.CloudTask) if task: self.assertEqual(_e, _m, task.exit_conditions.default.job_action.value, 'terminate') self.assertEqual(_e, _m, task.exit_conditions.exit_codes[0].code, 1) self.assertEqual(_e, _m, task.exit_conditions.exit_codes[0].exit_options.job_action.value, 'none') _m = "Test Create Task with Application Package" LOG.debug(_m) application = self.assertRuns(_e, _m, self.batch_mgmt_client.application.add_application, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id', allow_updated=True, display_name='my_application_name') package_ref = self.assertRuns(_e, _m, self.batch_mgmt_client.application.add_application_package, AZURE_RESOURCE_GROUP, AZURE_BATCH_ACCOUNT, 'my_application_id', 'v1.0') task = batch.models.TaskAddParameter('python_task_with_app_package', 'cmd /c echo hello world', application_package_references=[batch.models.ApplicationPackageReference('my_application_id')]) response = self.assertRuns(_e, _m, self.batch_client.task.add, 'python_test_job', task) self.assertIsNone(_e, _m, response) _m = "Test Get Task with Application Package" LOG.debug(_m) task = self.assertRuns(_e, _m, self.batch_client.task.get, 'python_test_job', 'python_task_with_app_package') self.assertIsInstance(task, batch.models.CloudTask) if task: self.assertEqual(_e, _m, task.id, 'python_task_with_app_package') self.assertEqual(_e, _m, task.application_package_references[0].application_id, 'my_application_id') response = self.assertRuns(_e, _m, self.batch_client.task.delete, 'python_test_job', 'python_task_with_app_package') self.assertIsNone(_e, _m, response) _m = "Test Create Task" LOG.debug(_m) task = batch.models.TaskAddParameter('python_task_1', 'cmd /c echo hello world') response = self.assertRuns(_e, _m, self.batch_client.task.add, 'python_test_job', task) self.assertIsNone(_e, _m, response) _m = "Test Terminate Task" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.task.terminate, 'python_test_job', 'python_task_1') self.assertIsNone(_e, _m, response) _m = "Test Update Task" LOG.debug(_m) task = batch.models.TaskAddParameter('python_task_2', 'cmd /c echo hello world') response = self.assertRuns(_e, _m, self.batch_client.task.add, 'python_test_job', task) constraints = batch.models.TaskConstraints(max_task_retry_count=1) response = self.assertRuns(_e, _m, self.batch_client.task.update, 'python_test_job', 'python_task_2', constraints=constraints) self.assertIsNone(_e, _m, response) _m = "Test Add Task Collection" LOG.debug(_m) tasks = [] for i in range(3, 6): tasks.append(batch.models.TaskAddParameter('python_task_{}'.format(i), 'cmd /c echo hello world')) response = self.assertRuns(_e, _m, self.batch_client.task.add_collection, 'python_test_job', tasks) self.assertIsInstance(response, batch.models.TaskAddCollectionResult) _m = "Test List Tasks" LOG.debug(_m) tasks = self.assertList(_e, _m, self.batch_client.task.list, 'python_test_job') self.assertEqual(_e, _m, len(tasks), 5) task_ids = [t.id for t in tasks] _m = "Test Get Task" LOG.debug(_m) task = self.assertRuns(_e, _m, self.batch_client.task.get, 'python_test_job', 'python_task_2') self.assertIsInstance(task, batch.models.CloudTask) if task: self.assertEqual(_e, _m, task.constraints.max_task_retry_count, 1) self.assertEqual(_e, _m, task.id, 'python_task_2') self.assertEqual(_e, _m, task.command_line, 'cmd /c echo hello world') _m = "Test Get Subtasks" #TODO: Test with actual subtasks LOG.debug('TODO: ' + _m) subtasks = self.assertRuns(_e, _m, self.batch_client.task.list_subtasks, 'python_test_job', 'python_task_2') #TODO: Assert LOG.debug(" wait for job to finish...") self.sleep(30) _m = "Test List Files" LOG.debug(_m) files = self.assertList(_e, _m, self.batch_client.file.list_from_task, 'python_test_job', 'python_task_2') self.assertTrue(_e, _m, len(files) > 2) _m = "Test Get File Properties" LOG.debug(_m) props = self.assertRuns(_e, _m, self.batch_client.file.get_node_file_properties_from_task, 'python_test_job', 'python_task_2', 'stdout.txt', raw=True) if props: self.assertTrue(_e, _m, 'Content-Length' in props.headers) self.assertTrue(_e, _m, 'Content-Type'in props.headers) _m = "Test Task File" LOG.debug(_m) file_length = 0 with io.BytesIO() as file_handle: response = self.assertRuns(_e, _m, self.batch_client.file.get_from_task, 'python_test_job', 'python_task_2', 'stdout.txt') if response: for data in response: file_length += len(data) self.assertTrue(_e, _m, file_length > 0) _m = "Test Prep and Release Task" #TODO: Test with actual prep and release tasks LOG.debug('TODO: ' + _m) response = self.assertRuns(_e, _m, self.batch_client.job.list_preparation_and_release_task_status, 'python_test_job') #TODO: Assert _m = "Test Delete File" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.file.delete_from_task, 'python_test_job', 'python_task_2', 'stdout.txt') self.assertIsNone(_e, _m, response) _m = "Test Delete Task" LOG.debug(_m) for id in task_ids: response = self.assertRuns(_e, _m, self.batch_client.task.delete, 'python_test_job', id) self.assertIsNone(_e, _m, response) _m = "Test Terminate Job" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.job.terminate, 'python_test_job') self.assertIsNone(_e, _m, response) _m = "Test Delete Job" LOG.debug(_m) for id in job_ids: response = self.assertRuns(_e, _m, self.batch_client.job.delete, id) self.assertIsNone(_e, _m, response) _m = "Test Job Lifetime Statistics" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.job.get_all_jobs_lifetime_statistics) self.assertSuccess(_e) @record def test_batch_job_schedules(self): _e = {} with BatchPool(self.live, self.batch_client, 'python_test_pool_5', target_dedicated=2) as pool_id: _m = "Test Create Job Schedule" LOG.debug(_m) job_spec = batch.models.JobSpecification(pool_info=batch.models.PoolInformation(pool_id)) schedule = batch.models.Schedule(start_window=datetime.timedelta(hours=1), recurrence_interval=datetime.timedelta(days=1)) params = batch.models.JobScheduleAddParameter('python_test_schedule', schedule, job_spec) response = self.assertRuns(_e, _m, self.batch_client.job_schedule.add, params) self.assertIsNone(_e, _m, response) _m = "Test List Job Schedules" LOG.debug(_m) schedules = self.assertList(_e, _m, self.batch_client.job_schedule.list) self.assertTrue(_e, _m, len(schedules) > 0) _m = "Test Get Job Schedule" LOG.debug(_m) schedule = self.assertRuns(_e, _m, self.batch_client.job_schedule.get, 'python_test_schedule') self.assertIsInstance(schedule, batch.models.CloudJobSchedule) if schedule: self.assertEqual(_e, _m, schedule.id, 'python_test_schedule') self.assertEqual(_e, _m, schedule.state, batch.models.JobScheduleState.active) _m = "Test Job Schedule Exists" LOG.debug(_m) exists = self.assertRuns(_e, _m, self.batch_client.job_schedule.exists, 'python_test_schedule') self.assertTrue(_e, _m, exists) _m = "Test List Jobs from Schedule" LOG.debug(_m) jobs = self.assertList(_e, _m, self.batch_client.job.list_from_job_schedule, 'python_test_schedule') self.assertTrue(_e, _m, len(jobs) > 0) if jobs: self.assertRuns(_e, _m, self.batch_client.job.delete, jobs[0].id) _m = "Test Disable Job Schedule" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.job_schedule.disable, 'python_test_schedule') self.assertIsNone(_e, _m, response) _m = "Test Enable Job Schedule" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.job_schedule.enable, 'python_test_schedule') self.assertIsNone(_e, _m, response) _m = "Test Update Job Schedule" LOG.debug(_m) job_spec = batch.models.JobSpecification(pool_info=batch.models.PoolInformation(pool_id)) schedule = batch.models.Schedule(recurrence_interval=datetime.timedelta(hours=10)) params = batch.models.JobScheduleUpdateParameter(schedule, job_spec) response = self.assertRuns(_e, _m, self.batch_client.job_schedule.update, 'python_test_schedule', params) self.assertIsNone(_e, _m, response) _m = "Test Patch Job Schedule" LOG.debug(_m) schedule = batch.models.Schedule(recurrence_interval=datetime.timedelta(hours=5)) params = batch.models.JobSchedulePatchParameter(schedule) response = self.assertRuns(_e, _m, self.batch_client.job_schedule.patch, 'python_test_schedule', params) self.assertIsNone(_e, _m, response) _m = "Test Terminate Job Schedule" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.job_schedule.terminate, 'python_test_schedule') self.assertIsNone(_e, _m, response) _m = "Test Delete Job Schedule" LOG.debug(_m) response = self.assertRuns(_e, _m, self.batch_client.job_schedule.delete, 'python_test_schedule') self.assertIsNone(_e, _m, response) self.assertSuccess(_e) class BatchPool(object): def __init__(self, live, client, id, **kwargs): self.live = live self.client = client self.id = id self.kwargs = kwargs self.nodes = kwargs.get('target_dedicated', 0) self.timeout = datetime.datetime.now() + datetime.timedelta(minutes=15) def __enter__(self): response = None if self.live: try: pool_config = batch.models.CloudServiceConfiguration('4') pool = batch.models.PoolAddParameter(self.id, 'small', cloud_service_configuration=pool_config, **self.kwargs) response = self.client.pool.add(pool) if self.live: LOG.debug(" waiting for pool to be ready...") time.sleep(30) ready_nodes = [] options = batch.models.ComputeNodeListOptions(select='id,state') while len(ready_nodes) != self.nodes: if datetime.datetime.now() > self.timeout: raise AssertionError("Nodes failed to become idle in 15 minutes") LOG.debug(" waiting for {} nodes to be ready...".format(self.nodes)) nodes = self.client.compute_node.list(self.id, options) ready_nodes = [n for n in nodes if n.state == batch.models.ComputeNodeState.idle] time.sleep(30) except Exception as err: try: self.client.pool.delete(self.id) except: pass raise AssertionError("Failed to create test pool: {}".format(err)) if response is None: return self.id else: raise AssertionError('Failed to create test pool') def __exit__(self, type, value, traceback): if self.live: try: self.client.pool.delete(self.id) except: pass if __name__ == '__main__': unittest.main()