diff --git a/airflow/providers/microsoft/azure/hooks/wasb.py b/airflow/providers/microsoft/azure/hooks/wasb.py index 44d6c569a1bbe..ad8cf2e7a0400 100644 --- a/airflow/providers/microsoft/azure/hooks/wasb.py +++ b/airflow/providers/microsoft/azure/hooks/wasb.py @@ -444,3 +444,14 @@ def delete_file( raise AirflowException(f'Blob(s) not found: {blob_name}') self.delete_blobs(container_name, *blobs_to_delete, **kwargs) + + def test_connection(self): + """Test Azure Blob Storage connection.""" + success = (True, "Successfully connected to Azure Blob Storage.") + + try: + # Attempt to retrieve storage account information + self.get_conn().get_account_information() + return success + except Exception as e: + return False, str(e) diff --git a/tests/providers/microsoft/azure/hooks/test_wasb.py b/tests/providers/microsoft/azure/hooks/test_wasb.py index e31a560ff0821..c4a31d671bb22 100644 --- a/tests/providers/microsoft/azure/hooks/test_wasb.py +++ b/tests/providers/microsoft/azure/hooks/test_wasb.py @@ -369,3 +369,25 @@ def test_delete_multiple_nonexisting_blobs_fails(self, mock_getblobs): hook = WasbHook(wasb_conn_id=self.shared_key_conn_id) hook.delete_file('container', 'nonexisting_blob_prefix', is_prefix=True, ignore_if_missing=False) assert isinstance(ctx.value, AirflowException) + + @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.BlobServiceClient") + def test_connection_success(self, mock_service): + hook = WasbHook(wasb_conn_id=self.shared_key_conn_id) + hook.get_conn().get_account_information().return_value = { + 'sku_name': 'Standard_RAGRS', + 'account_kind': 'StorageV2', + } + status, msg = hook.test_connection() + + assert status is True + assert msg == "Successfully connected to Azure Blob Storage." + + @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.BlobServiceClient") + def test_connection_failure(self, mock_service): + hook = WasbHook(wasb_conn_id=self.shared_key_conn_id) + hook.get_conn().get_account_information = mock.PropertyMock( + side_effect=Exception("Authentication failed.") + ) + status, msg = hook.test_connection() + assert status is False + assert msg == "Authentication failed."