Add test_connection method to Databricks hook (#24617)

This PR enables the test button on the airflow UI for testing whether the Databricks connection works using the connection params filled by the user

airflow/providers/databricks/hooks/databricks.py

@@ -52,6 +52,8 @@
 
 RUN_LIFE_CYCLE_STATES = ['PENDING', 'RUNNING', 'TERMINATING', 'TERMINATED', 'SKIPPED', 'INTERNAL_ERROR']
 
+LIST_ZONES_ENDPOINT = ('GET', 'api/2.0/clusters/list-zones')
+
 
 class RunState:
     """Utility class for the run state concept of Databricks runs."""
@@ -408,3 +410,16 @@ def get_repo_by_path(self, path: str) -> Optional[str]:
                 raise e
 
         return None
+
+    def test_connection(self):
+        """Test the Databricks connectivity from UI"""
+        hook = DatabricksHook(databricks_conn_id=self.databricks_conn_id)
+        try:
+            hook._do_api_call(endpoint_info=LIST_ZONES_ENDPOINT).get('zones')
+            status = True
+            message = 'Connection successfully tested'
+        except Exception as e:
+            status = False
+            message = str(e)
+
+        return status, message

tests/providers/databricks/hooks/test_databricks.py

@@ -104,6 +104,7 @@
     ],
     'has_more': False,
 }
+LIST_ZONES_RESPONSE = {'zones': ['us-east-2b', 'us-east-2c', 'us-east-2a'], 'default_zone': 'us-east-2b'}
 
 
 def run_now_endpoint(host):
@@ -178,11 +179,16 @@ def uninstall_endpoint(host):
 
 def list_jobs_endpoint(host):
     """
-    Utility function to generate the list jobs endpoint giver the host
+    Utility function to generate the list jobs endpoint given the host
     """
     return f'https://{host}/api/2.1/jobs/list'
 
 
+def list_zones_endpoint(host):
+    """Utility function to generate the list zones endpoint given the host"""
+    return f'https://{host}/api/2.0/clusters/list-zones'
+
+
 def create_valid_response_mock(content):
     response = mock.MagicMock()
     response.json.return_value = content
@@ -706,6 +712,40 @@ def test_get_job_id_by_name_raise_exception_with_duplicates(self, mock_requests)
             timeout=self.hook.timeout_seconds,
         )
 
+    @mock.patch('airflow.providers.databricks.hooks.databricks_base.requests')
+    def test_connection_success(self, mock_requests):
+        mock_requests.codes.ok = 200
+        mock_requests.get.return_value.json.return_value = LIST_ZONES_RESPONSE
+        status_code_mock = mock.PropertyMock(return_value=200)
+        type(mock_requests.get.return_value).status_code = status_code_mock
+        response = self.hook.test_connection()
+        assert response == (True, 'Connection successfully tested')
+        mock_requests.get.assert_called_once_with(
+            list_zones_endpoint(HOST),
+            json=None,
+            params=None,
+            auth=HTTPBasicAuth(LOGIN, PASSWORD),
+            headers=USER_AGENT_HEADER,
+            timeout=self.hook.timeout_seconds,
+        )
+
+    @mock.patch('airflow.providers.databricks.hooks.databricks_base.requests')
+    def test_connection_failure(self, mock_requests):
+        mock_requests.codes.ok = 404
+        mock_requests.get.side_effect = Exception('Connection Failure')
+        status_code_mock = mock.PropertyMock(return_value=404)
+        type(mock_requests.get.return_value).status_code = status_code_mock
+        response = self.hook.test_connection()
+        assert response == (False, 'Connection Failure')
+        mock_requests.get.assert_called_once_with(
+            list_zones_endpoint(HOST),
+            json=None,
+            params=None,
+            auth=HTTPBasicAuth(LOGIN, PASSWORD),
+            headers=USER_AGENT_HEADER,
+            timeout=self.hook.timeout_seconds,
+        )
+
 
 class TestDatabricksHookToken(unittest.TestCase):
     """