From 86b0e083c8758ba6744c9ccbcf64a4010c1fd995 Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Fri, 22 Jul 2022 19:37:23 +0530 Subject: [PATCH 1/2] Add test_connection to Azure Batch hook --- .../providers/microsoft/azure/hooks/batch.py | 13 +++++++++++++ .../microsoft/azure/hooks/test_azure_batch.py | 18 ++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/airflow/providers/microsoft/azure/hooks/batch.py b/airflow/providers/microsoft/azure/hooks/batch.py index 7e1f0ac5bb749..eb4a48ac18e66 100644 --- a/airflow/providers/microsoft/azure/hooks/batch.py +++ b/airflow/providers/microsoft/azure/hooks/batch.py @@ -362,3 +362,16 @@ def wait_for_job_tasks_to_complete(self, job_id: str, timeout: int) -> None: self.log.info("Waiting for %s to complete, currently on %s state", task.id, task.state) time.sleep(15) raise TimeoutError("Timed out waiting for tasks to complete") + + def test_connection(self): + """Test a configured Azure Batch connection.""" + try: + # Attempt to list existing jobs under the configured Batch account and retrieve + # the first in the returned iterator. The Azure Batch API does allow for creation of a + # BatchServiceClient with incorrect values but then will fail properly once items are + # retrieved using the client. We need to _actually_ try to retrieve an object to properly + # test the connection. + next(self.get_conn().job.list(), None) + except Exception as e: + return False, str(e) + return True, "Successfully connected to Azure Batch." diff --git a/tests/providers/microsoft/azure/hooks/test_azure_batch.py b/tests/providers/microsoft/azure/hooks/test_azure_batch.py index daed87c30ea9e..99d7f68166caf 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_batch.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_batch.py @@ -19,6 +19,7 @@ import json import unittest from unittest import mock +from unittest.mock import PropertyMock from azure.batch import BatchServiceClient, models as batch_models @@ -165,3 +166,20 @@ def test_add_single_task_to_job(self, mock_batch): def test_wait_for_all_task_to_complete(self, mock_batch): # TODO: Add test pass + + @mock.patch('airflow.providers.microsoft.azure.hooks.batch.BatchServiceClient') + def test_connection_success(self, mock_batch): + hook = AzureBatchHook(azure_batch_conn_id=self.test_cloud_conn_id) + hook.get_conn().job.return_value = {} + status, msg = hook.test_connection() + print(status, msg) + assert status is True + assert msg == "Successfully connected to Azure Batch." + + @mock.patch('airflow.providers.microsoft.azure.hooks.batch.BatchServiceClient') + def test_connection_failure(self, mock_batch): + hook = AzureBatchHook(azure_batch_conn_id=self.test_cloud_conn_id) + hook.get_conn().job.list = PropertyMock(side_effect=Exception("Authentication failed.")) + status, msg = hook.test_connection() + assert status is False + assert msg == "Authentication failed." From eb5a00846d35b5d1fc4e493bb9500f6d7487b274 Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Fri, 22 Jul 2022 19:58:42 +0530 Subject: [PATCH 2/2] Apply review suggestions --- tests/providers/microsoft/azure/hooks/test_azure_batch.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/providers/microsoft/azure/hooks/test_azure_batch.py b/tests/providers/microsoft/azure/hooks/test_azure_batch.py index 99d7f68166caf..94e4d1bc560ed 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_batch.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_batch.py @@ -172,7 +172,6 @@ def test_connection_success(self, mock_batch): hook = AzureBatchHook(azure_batch_conn_id=self.test_cloud_conn_id) hook.get_conn().job.return_value = {} status, msg = hook.test_connection() - print(status, msg) assert status is True assert msg == "Successfully connected to Azure Batch."