def test_job_configuration_and_create_job(self, mock_batch): hook = AzureBatchHook(azure_batch_conn_id=self.test_vm_conn_id) mock_instance = mock_batch.return_value.job.add job = hook.configure_job(job_id='myjob', pool_id='mypool') hook.create_job(job) self.assertIsInstance(job, batch_models.JobAddParameter) mock_instance.assert_called_once_with(job)
def test_add_single_task_to_job(self, mock_batch): hook = AzureBatchHook(azure_batch_conn_id=self.test_vm_conn_id) mock_instance = mock_batch.return_value.task.add task = hook.configure_task(task_id="mytask", command_line="echo hello") hook.add_single_task_to_job(job_id='myjob', task=task) self.assertIsInstance(task, batch_models.TaskAddParameter) mock_instance.assert_called_once_with(job_id="myjob", task=task)
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)
def get_hook(self) -> AzureBatchHook: """ Create and return an AzureBatchHook. """ return AzureBatchHook( azure_batch_conn_id=self.azure_batch_conn_id )
def test_connection_and_client(self): hook = AzureBatchHook(azure_batch_conn_id=self.test_vm_conn_id) self.assertIsInstance(hook._connection(), Connection) self.assertIsInstance(hook.get_conn(), BatchServiceClient)