From 5bf1cd06d1bbcffbd62f8e59db40989e936f282c Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Wed, 27 Jul 2022 17:29:15 +0530 Subject: [PATCH 1/3] Add test_connection method to AzureContainerInstanceHook --- .../microsoft/azure/hooks/container_instance.py | 12 ++++++++++++ .../azure/hooks/test_azure_container_instance.py | 14 ++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/airflow/providers/microsoft/azure/hooks/container_instance.py b/airflow/providers/microsoft/azure/hooks/container_instance.py index 9b0cd5d17264f..e185fcfb1ea74 100644 --- a/airflow/providers/microsoft/azure/hooks/container_instance.py +++ b/airflow/providers/microsoft/azure/hooks/container_instance.py @@ -138,3 +138,15 @@ def exists(self, resource_group: str, name: str) -> bool: 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/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py b/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py index e29d9c2881e15..884bc3fd7b478 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py @@ -97,3 +97,17 @@ def test_exists_with_not_existing(self, list_mock): ) ] 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." From 6f3a8d2786321f7203aa3fd3424c3e2831b1031f Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Thu, 28 Jul 2022 16:56:59 +0530 Subject: [PATCH 2/3] Use conn_name_attr in the init method --- .../providers/microsoft/azure/hooks/container_instance.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/container_instance.py b/airflow/providers/microsoft/azure/hooks/container_instance.py index e185fcfb1ea74..2cc391b347d0c 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` of + :param azure_conn_id: :ref:`Azure connection id` 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: From b6f9a1d7795b69fb6d6fce57e83dab766293f429 Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Thu, 28 Jul 2022 17:19:16 +0530 Subject: [PATCH 3/3] Use azure_conn_id instead of conn_id --- .../providers/microsoft/azure/operators/container_instances.py | 2 +- .../microsoft/azure/hooks/test_azure_container_instance.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/container_instances.py b/airflow/providers/microsoft/azure/operators/container_instances.py index 519ce8fe41937..2a33f89fd54d0 100644 --- a/airflow/providers/microsoft/azure/operators/container_instances.py +++ b/airflow/providers/microsoft/azure/operators/container_instances.py @@ -184,7 +184,7 @@ def execute(self, context: "Context") -> int: # 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 884bc3fd7b478..289e9823b8c60 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 @@ def setUp(self): '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')