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 dc9b516494 Adds connection test for ADLS Gen2 (#32126)
dc9b516494 is described below
commit dc9b516494ad5587d30b19d3b7cffc198c27a52c
Author: Akash Sharma <[email protected]>
AuthorDate: Mon Jun 26 22:32:10 2023 +0530
Adds connection test for ADLS Gen2 (#32126)
---
.../providers/microsoft/azure/hooks/data_lake.py | 13 +++++++++++
.../microsoft/azure/hooks/test_azure_data_lake.py | 26 ++++++++++++++++++++++
2 files changed, 39 insertions(+)
diff --git a/airflow/providers/microsoft/azure/hooks/data_lake.py
b/airflow/providers/microsoft/azure/hooks/data_lake.py
index 1fb6c050be..5098113baa 100644
--- a/airflow/providers/microsoft/azure/hooks/data_lake.py
+++ b/airflow/providers/microsoft/azure/hooks/data_lake.py
@@ -509,3 +509,16 @@ class AzureDataLakeStorageV2Hook(BaseHook):
"""
directory_client = self.get_directory_client(file_system_name,
directory_name)
directory_client.delete_directory()
+
+ def test_connection(self):
+ """Test ADLS Gen2 Storage connection."""
+ try:
+ # Attempts to list file systems in ADLS Gen2 Storage and retrieves
the first
+ # file_system from the returned iterator. The Azure DataLake
Storage allows creation
+ # of DataLakeServiceClient even if the credentials are incorrect
but will fail properly
+ # if we try to fetch the file_system. We need to _actually_ try to
retrieve a
+ # file_system to properly test the connection
+ next(self.get_conn().list_file_systems(), None)
+ return True, "Successfully connected to ADLS Gen2 Storage."
+ except Exception as e:
+ return False, str(e)
diff --git a/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py
b/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py
index e8d378ab5c..161b8d9995 100644
--- a/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py
+++ b/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py
@@ -19,6 +19,10 @@ from __future__ import annotations
import json
from unittest import mock
+from unittest.mock import PropertyMock
+
+import pytest
+from azure.storage.filedatalake._models import FileSystemProperties
from airflow.models import Connection
from airflow.providers.microsoft.azure.hooks.data_lake import
AzureDataLakeHook, AzureDataLakeStorageV2Hook
@@ -245,3 +249,25 @@ class TestAzureDataLakeStorageV2Hook:
hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id)
hook.list_files_directory(self.file_system_name, self.directory_name)
mock_get_file_system.return_value.get_paths.assert_called_once_with(self.directory_name)
+
+ @pytest.mark.parametrize(
+ argnames="list_file_systems_result",
+ argvalues=[iter([FileSystemProperties]), iter([])],
+ )
+
@mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn")
+ def test_connection_success(self, mock_conn, list_file_systems_result):
+ hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id)
+ hook.get_conn().list_file_systems.return_value =
list_file_systems_result
+ status, msg = hook.test_connection()
+
+ assert status is True
+ assert msg == "Successfully connected to ADLS Gen2 Storage."
+
+
@mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn")
+ def test_connection_failure(self, mock_conn):
+ hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id)
+ hook.get_conn().list_file_systems =
PropertyMock(side_effect=Exception("Authentication failed."))
+ status, msg = hook.test_connection()
+
+ assert status is False
+ assert msg == "Authentication failed."
