def test_configure_pool_with_cloud_config(self):
hook = AzureBatchHook(azure_batch_conn_id=self.test_cloud_conn_id)
pool = hook.configure_pool(
pool_id='mypool',
vm_size="test_vm_size",
target_dedicated_nodes=1,
)
self.assertIsInstance(pool, batch_models.PoolAddParameter)
def test_create_pool_with_cloud_config(self, mock_batch):
hook = AzureBatchHook(azure_batch_conn_id=self.test_cloud_conn_id)
mock_instance = mock_batch.return_value.pool.add
pool = hook.configure_pool(
pool_id='mypool',
vm_size="test_vm_size",
target_dedicated_nodes=1,
)
hook.create_pool(pool=pool)
mock_instance.assert_called_once_with(pool)
def test_configure_pool_with_cloud_config(self):
hook = AzureBatchHook(azure_batch_conn_id=self.test_cloud_conn_id)
pool = hook.configure_pool(
pool_id='mypool',
vm_size="test_vm_size",
target_dedicated_nodes=1,
vm_publisher="test.vm.publisher",
vm_offer="test.vm.offer",
sku_starts_with="test-sku",
)
self.assertIsInstance(pool, batch_models.PoolAddParameter)
def test_create_pool_with_vm_config(self, mock_batch):
hook = AzureBatchHook(azure_batch_conn_id=self.test_vm_conn_id)
mock_instance = mock_batch.return_value.pool.add
pool = hook.configure_pool(
pool_id='mypool',
vm_size="test_vm_size",
target_dedicated_nodes=1,
vm_publisher="test.vm.publisher",
vm_offer="test.vm.offer",
sku_starts_with="test-sku",
)
hook.create_pool(pool=pool)
mock_instance.assert_called_once_with(pool)
def test_configure_pool_with_latest_vm(self):
with mock.patch("airflow.providers.microsoft.azure.hooks."
"azure_batch.AzureBatchHook._get_latest_verified_image_vm_and_sku")\
as mock_getvm:
hook = AzureBatchHook(azure_batch_conn_id=self.test_cloud_conn_id)
getvm_instance = mock_getvm
getvm_instance.return_value = ['test-image', 'test-sku']
pool = hook.configure_pool(
pool_id='mypool',
vm_size="test_vm_size",
use_latest_image_and_sku=True,
)
self.assertIsInstance(pool, batch_models.PoolAddParameter)