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 b27fc0367c Add test_connection method to AzureFileShareHook (#24843)
b27fc0367c is described below
commit b27fc0367cd1125f4d312497ba5337115476315e
Author: Phani Kumar <[email protected]>
AuthorDate: Wed Jul 6 19:18:15 2022 +0530
Add test_connection method to AzureFileShareHook (#24843)
---
airflow/providers/microsoft/azure/hooks/fileshare.py | 16 ++++++++++++++++
.../microsoft/azure/hooks/test_azure_fileshare.py | 16 ++++++++++++++++
2 files changed, 32 insertions(+)
diff --git a/airflow/providers/microsoft/azure/hooks/fileshare.py
b/airflow/providers/microsoft/azure/hooks/fileshare.py
index 1f78d6e4b2..1262746c64 100644
--- a/airflow/providers/microsoft/azure/hooks/fileshare.py
+++ b/airflow/providers/microsoft/azure/hooks/fileshare.py
@@ -286,3 +286,19 @@ class AzureFileShareHook(BaseHook):
self.get_conn().create_file_from_stream(
share_name, directory_name, file_name, stream, count, **kwargs
)
+
+ def test_connection(self):
+ """Test Azure FileShare connection."""
+ success = (True, "Successfully connected to Azure File Share.")
+
+ try:
+ # Attempt to retrieve file share information
+ next(iter(self.get_conn().list_shares()))
+ return success
+ except StopIteration:
+ # If the iterator returned is empty it should still be considered
a successful connection since
+ # it's possible to create a storage account without any file share
and none could
+ # legitimately exist yet.
+ return success
+ except Exception as e:
+ return False, str(e)
diff --git a/tests/providers/microsoft/azure/hooks/test_azure_fileshare.py
b/tests/providers/microsoft/azure/hooks/test_azure_fileshare.py
index 57c6faf824..8fda79a68a 100644
--- a/tests/providers/microsoft/azure/hooks/test_azure_fileshare.py
+++ b/tests/providers/microsoft/azure/hooks/test_azure_fileshare.py
@@ -255,3 +255,19 @@ class TestAzureFileshareHook(unittest.TestCase):
hook =
AzureFileShareHook(azure_fileshare_conn_id='azure_fileshare_extras')
hook.delete_share('my_share')
mock_instance.delete_share.assert_called_once_with('my_share')
+
+
@mock.patch('airflow.providers.microsoft.azure.hooks.fileshare.FileService',
autospec=True)
+ def test_connection_success(self, mock_service):
+ hook =
AzureFileShareHook(azure_fileshare_conn_id='azure_fileshare_extras')
+ hook.get_conn().list_shares.return_value = ["test_container"]
+ status, msg = hook.test_connection()
+ assert status is True
+ assert msg == "Successfully connected to Azure File Share."
+
+
@mock.patch('airflow.providers.microsoft.azure.hooks.fileshare.FileService',
autospec=True)
+ def test_connection_failure(self, mock_service):
+ hook =
AzureFileShareHook(azure_fileshare_conn_id='azure_fileshare_extras')
+ hook.get_conn().list_shares.side_effect = Exception("Test Connection
Failure")
+ status, msg = hook.test_connection()
+ assert status is False
+ assert msg == "Test Connection Failure"
