[airflow] branch main updated: Add `test_connection` method to AzureContainerInstanceHook (#25362)

[po...@apache.org]

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

potiuk 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 8bb0c4fd32 Add `test_connection` method to AzureContainerInstanceHook (#25362)
8bb0c4fd32 is described below

commit 8bb0c4fd32b21bf2900e33ec29b1dc7d772589c9
Author: Phani Kumar <94...@users.noreply.github.com>
AuthorDate: Thu Jul 28 19:10:03 2022 +0530

    Add `test_connection` method to AzureContainerInstanceHook (#25362)
---
 .../microsoft/azure/hooks/container_instance.py        | 18 +++++++++++++++---
 .../microsoft/azure/operators/container_instances.py   |  2 +-
 .../azure/hooks/test_azure_container_instance.py       | 16 +++++++++++++++-
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/airflow/providers/microsoft/azure/hooks/container_instance.py b/airflow/providers/microsoft/azure/hooks/container_instance.py
index 9b0cd5d172..2cc391b347 100644
--- a/airflow/providers/microsoft/azure/hooks/container_instance.py
+++ b/airflow/providers/microsoft/azure/hooks/container_instance.py
@@ -36,7 +36,7 @@ class AzureContainerInstanceHook(AzureBaseHook):
     client_id (Application ID) as login, the generated password as password,
     and tenantId and subscriptionId in the extra's field as a json.
 
-    :param conn_id: :ref:`Azure connection id<howto/connection:azure>` of
+    :param azure_conn_id: :ref:`Azure connection id<howto/connection:azure>` of
         a service principal which will be used to start the container instance.
     """
 
@@ -45,8 +45,8 @@ class AzureContainerInstanceHook(AzureBaseHook):
     conn_type = 'azure_container_instance'
     hook_name = 'Azure Container Instance'
 
-    def __init__(self, conn_id: str = default_conn_name) -> None:
-        super().__init__(sdk_client=ContainerInstanceManagementClient, conn_id=conn_id)
+    def __init__(self, azure_conn_id: str = default_conn_name) -> None:
+        super().__init__(sdk_client=ContainerInstanceManagementClient, conn_id=azure_conn_id)
         self.connection = self.get_conn()
 
     def create_or_update(self, resource_group: str, name: str, container_group: ContainerGroup) -> None:
@@ -138,3 +138,15 @@ class AzureContainerInstanceHook(AzureBaseHook):
             if container.name == name:
                 return True
         return False
+
+    def test_connection(self):
+        """Test a configured Azure Container Instance connection."""
+        try:
+            # Attempt to list existing container groups under the configured subscription and retrieve the
+            # first in the returned iterator. We need to _actually_ try to retrieve an object to properly
+            # test the connection.
+            next(self.connection.container_groups.list(), None)
+        except Exception as e:
+            return False, str(e)
+
+        return True, "Successfully connected to Azure Container Instance."
diff --git a/airflow/providers/microsoft/azure/operators/container_instances.py b/airflow/providers/microsoft/azure/operators/container_instances.py
index 519ce8fe41..2a33f89fd5 100644
--- a/airflow/providers/microsoft/azure/operators/container_instances.py
+++ b/airflow/providers/microsoft/azure/operators/container_instances.py
@@ -184,7 +184,7 @@ class AzureContainerInstancesOperator(BaseOperator):
         # Check name again in case it was templated.
         self._check_name(self.name)
 
-        self._ci_hook = AzureContainerInstanceHook(conn_id=self.ci_conn_id)
+        self._ci_hook = AzureContainerInstanceHook(azure_conn_id=self.ci_conn_id)
 
         if self.fail_if_exists:
             self.log.info("Testing if container group already exists")
diff --git a/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py b/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py
index e29d9c2881..289e9823b8 100644
--- a/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py
+++ b/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py
@@ -50,7 +50,7 @@ class TestAzureContainerInstanceHook(unittest.TestCase):
             'azure.common.credentials.ServicePrincipalCredentials.__init__', autospec=True, return_value=None
         ):
             with patch('azure.mgmt.containerinstance.ContainerInstanceManagementClient'):
-                self.hook = AzureContainerInstanceHook(conn_id='azure_container_instance_test')
+                self.hook = AzureContainerInstanceHook(azure_conn_id='azure_container_instance_test')
 
     @patch('azure.mgmt.containerinstance.models.ContainerGroup')
     @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.create_or_update')
@@ -97,3 +97,17 @@ class TestAzureContainerInstanceHook(unittest.TestCase):
             )
         ]
         assert not self.hook.exists('test', 'not found')
+
+    @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.list')
+    def test_connection_success(self, mock_container_groups_list):
+        mock_container_groups_list.return_value = iter([])
+        status, msg = self.hook.test_connection()
+        assert status is True
+        assert msg == "Successfully connected to Azure Container Instance."
+
+    @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.list')
+    def test_connection_failure(self, mock_container_groups_list):
+        mock_container_groups_list.side_effect = Exception("Authentication failed.")
+        status, msg = self.hook.test_connection()
+        assert status is False
+        assert msg == "Authentication failed."