From fd682fd70a97a1f937786a1a136f0fa929c8fb80 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 5 Oct 2020 08:39:50 +0100 Subject: [PATCH] fix job deletion (#11272) --- .../providers/microsoft/azure/operators/azure_batch.py | 2 +- .../microsoft/azure/operators/test_azure_batch.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/operators/azure_batch.py b/airflow/providers/microsoft/azure/operators/azure_batch.py index 54b8e36015..5de640eaeb 100644 --- a/airflow/providers/microsoft/azure/operators/azure_batch.py +++ b/airflow/providers/microsoft/azure/operators/azure_batch.py @@ -350,7 +350,7 @@ class AzureBatchOperator(BaseOperator): """ if 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: self.log.info("Deleting pool: %s", pool_id) self.hook.connection.pool.delete(pool_id) diff --git a/tests/providers/microsoft/azure/operators/test_azure_batch.py b/tests/providers/microsoft/azure/operators/test_azure_batch.py index 89ed85c208..5f7132cdeb 100644 --- a/tests/providers/microsoft/azure/operators/test_azure_batch.py +++ b/tests/providers/microsoft/azure/operators/test_azure_batch.py @@ -138,6 +138,7 @@ class TestAzureBatchOperator(unittest.TestCase): # pylint: disable=too-many-ins timeout=2, ) self.batch_client = mock_batch.return_value + self.mock_instance = mock_hook.return_value self.assertEqual(self.batch_client, self.operator.hook.connection) @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: self.operator2_no_formula.execute(None) 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')