[airflow] branch main updated: Add test_connection to Azure Batch hook (#25235)

[jo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

joshfell pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new eab0167f1b Add test_connection to Azure Batch hook (#25235)
eab0167f1b is described below

commit eab0167f1beb81de8e613685da79ef9a04eef5b3
Author: Phani Kumar <94...@users.noreply.github.com>
AuthorDate: Sat Jul 23 02:36:40 2022 +0530

    Add test_connection to Azure Batch hook (#25235)
    
    * Add test_connection to Azure Batch hook
    
    * Apply review suggestions
---
 airflow/providers/microsoft/azure/hooks/batch.py        | 13 +++++++++++++
 .../providers/microsoft/azure/hooks/test_azure_batch.py | 17 +++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/airflow/providers/microsoft/azure/hooks/batch.py b/airflow/providers/microsoft/azure/hooks/batch.py
index 7e1f0ac5bb..eb4a48ac18 100644
--- a/airflow/providers/microsoft/azure/hooks/batch.py
+++ b/airflow/providers/microsoft/azure/hooks/batch.py
@@ -362,3 +362,16 @@ class AzureBatchHook(BaseHook):
                 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 daed87c30e..94e4d1bc56 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,19 @@ class TestAzureBatchHook(unittest.TestCase):
     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()
+        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."