This commit is contained in:
Ephraim Anierobi 2020-10-05 08:39:50 +01:00 коммит произвёл GitHub
Родитель 1180f1ba42
Коммит fd682fd70a
Не найден ключ, соответствующий данной подписи
Идентификатор ключа GPG: 4AEE18F83AFDEB23
2 изменённых файлов: 11 добавлений и 1 удалений

Просмотреть файл

@ -350,7 +350,7 @@ class AzureBatchOperator(BaseOperator):
""" """
if job_id: if job_id:
self.log.info("Deleting job: %s", job_id) self.log.info("Deleting job: %s", job_id)
self.hook.connection.job.delete(pool_id) self.hook.connection.job.delete(job_id)
if pool_id: if pool_id:
self.log.info("Deleting pool: %s", pool_id) self.log.info("Deleting pool: %s", pool_id)
self.hook.connection.pool.delete(pool_id) self.hook.connection.pool.delete(pool_id)

Просмотреть файл

@ -138,6 +138,7 @@ class TestAzureBatchOperator(unittest.TestCase): # pylint: disable=too-many-ins
timeout=2, timeout=2,
) )
self.batch_client = mock_batch.return_value self.batch_client = mock_batch.return_value
self.mock_instance = mock_hook.return_value
self.assertEqual(self.batch_client, self.operator.hook.connection) self.assertEqual(self.batch_client, self.operator.hook.connection)
@mock.patch.object(AzureBatchHook, 'wait_for_all_node_state') @mock.patch.object(AzureBatchHook, 'wait_for_all_node_state')
@ -194,3 +195,12 @@ class TestAzureBatchOperator(unittest.TestCase): # pylint: disable=too-many-ins
with self.assertRaises(AirflowException) as e: with self.assertRaises(AirflowException) as e:
self.operator2_no_formula.execute(None) self.operator2_no_formula.execute(None)
self.assertEqual(str(e.exception), "The auto_scale_formula is required when enable_auto_scale is set") self.assertEqual(str(e.exception), "The auto_scale_formula is required when enable_auto_scale is set")
def test_cleaning_works(self):
self.operator.clean_up(job_id="myjob")
self.batch_client.job.delete.assert_called_once_with("myjob")
self.operator.clean_up("mypool")
self.batch_client.pool.delete.assert_called_once_with("mypool")
self.operator.clean_up("mypool", "myjob")
self.batch_client.job.delete.assert_called_with("myjob")
self.batch_client.pool.delete.assert_called_with('mypool')