You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2022/07/18 12:45:24 UTC
[airflow] branch main updated: Databricks: fix test_connection implementation (#25114)
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 2f70daf5ac Databricks: fix test_connection implementation (#25114)
2f70daf5ac is described below
commit 2f70daf5ac36100ff0bbd4ac66ce921a2bc6dea0
Author: Alex Ott <al...@gmail.com>
AuthorDate: Mon Jul 18 14:45:11 2022 +0200
Databricks: fix test_connection implementation (#25114)
Original implementation has used the API that exists only on AWS, so it will be failing on
Azure & GCP Databricks. This PR uses another API that is available on all Databricks
cloud platforms.
---
airflow/providers/databricks/hooks/databricks.py | 4 ++--
tests/providers/databricks/hooks/test_databricks.py | 18 +++++++++++-------
2 files changed, 13 insertions(+), 9 deletions(-)
diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py
index 33b4aeb8f7..9ae061b02e 100644
--- a/airflow/providers/databricks/hooks/databricks.py
+++ b/airflow/providers/databricks/hooks/databricks.py
@@ -52,7 +52,7 @@ WORKSPACE_GET_STATUS_ENDPOINT = ('GET', 'api/2.0/workspace/get-status')
RUN_LIFE_CYCLE_STATES = ['PENDING', 'RUNNING', 'TERMINATING', 'TERMINATED', 'SKIPPED', 'INTERNAL_ERROR']
-LIST_ZONES_ENDPOINT = ('GET', 'api/2.0/clusters/list-zones')
+SPARK_VERSIONS_ENDPOINT = ('GET', 'api/2.0/clusters/spark-versions')
class RunState:
@@ -415,7 +415,7 @@ class DatabricksHook(BaseDatabricksHook):
"""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')
+ hook._do_api_call(endpoint_info=SPARK_VERSIONS_ENDPOINT).get('versions')
status = True
message = 'Connection successfully tested'
except Exception as e:
diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py
index 8599ac9918..fa2c416a10 100644
--- a/tests/providers/databricks/hooks/test_databricks.py
+++ b/tests/providers/databricks/hooks/test_databricks.py
@@ -104,7 +104,11 @@ LIST_JOBS_RESPONSE = {
],
'has_more': False,
}
-LIST_ZONES_RESPONSE = {'zones': ['us-east-2b', 'us-east-2c', 'us-east-2a'], 'default_zone': 'us-east-2b'}
+LIST_SPARK_VERSIONS_RESPONSE = {
+ "versions": [
+ {"key": "8.2.x-scala2.12", "name": "8.2 (includes Apache Spark 3.1.1, Scala 2.12)"},
+ ]
+}
def run_now_endpoint(host):
@@ -184,9 +188,9 @@ def list_jobs_endpoint(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 list_spark_versions_endpoint(host):
+ """Utility function to generate the list spark versions endpoint given the host"""
+ return f'https://{host}/api/2.0/clusters/spark-versions'
def create_valid_response_mock(content):
@@ -715,13 +719,13 @@ class TestDatabricksHook(unittest.TestCase):
@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
+ mock_requests.get.return_value.json.return_value = LIST_SPARK_VERSIONS_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),
+ list_spark_versions_endpoint(HOST),
json=None,
params=None,
auth=HTTPBasicAuth(LOGIN, PASSWORD),
@@ -738,7 +742,7 @@ class TestDatabricksHook(unittest.TestCase):
response = self.hook.test_connection()
assert response == (False, 'Connection Failure')
mock_requests.get.assert_called_once_with(
- list_zones_endpoint(HOST),
+ list_spark_versions_endpoint(HOST),
json=None,
params=None,
auth=HTTPBasicAuth(LOGIN, PASSWORD),