From a7acca05e598be13d2dc602a6c7f776b808493b8 Mon Sep 17 00:00:00 2001 From: Hua Shi Date: Wed, 21 Aug 2024 22:12:52 -0700 Subject: [PATCH 1/3] base_url by default is empty string, should call when it takes the default value --- airflow/providers/apache/livy/hooks/livy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/apache/livy/hooks/livy.py b/airflow/providers/apache/livy/hooks/livy.py index 53a86b5cfb198..f35c9336c354e 100644 --- a/airflow/providers/apache/livy/hooks/livy.py +++ b/airflow/providers/apache/livy/hooks/livy.py @@ -155,7 +155,7 @@ def post_batch(self, *args: Any, **kwargs: Any) -> int: """ batch_submit_body = json.dumps(self.build_post_batch_body(*args, **kwargs)) - if self.base_url is None: + if self.base_url in (None, ""): # need to init self.base_url self.get_conn() self.log.info("Submitting job %s to %s", batch_submit_body, self.base_url) From c25fe63d63fe0d353d22f0ecaab4cb6a7fc7d15b Mon Sep 17 00:00:00 2001 From: "harry.shi" <38423541+harryshi10@users.noreply.github.com> Date: Wed, 28 Aug 2024 00:07:30 +0300 Subject: [PATCH 2/3] check base_url with if not Co-authored-by: Jarek Potiuk --- airflow/providers/apache/livy/hooks/livy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/apache/livy/hooks/livy.py b/airflow/providers/apache/livy/hooks/livy.py index f35c9336c354e..7cf6184aea15c 100644 --- a/airflow/providers/apache/livy/hooks/livy.py +++ b/airflow/providers/apache/livy/hooks/livy.py @@ -155,7 +155,7 @@ def post_batch(self, *args: Any, **kwargs: Any) -> int: """ batch_submit_body = json.dumps(self.build_post_batch_body(*args, **kwargs)) - if self.base_url in (None, ""): + if not self.base_url: # need to init self.base_url self.get_conn() self.log.info("Submitting job %s to %s", batch_submit_body, self.base_url) From 631658db3338c6193b16efc049d5e4353d8160e5 Mon Sep 17 00:00:00 2001 From: Hua Shi Date: Mon, 2 Sep 2024 10:08:07 -0700 Subject: [PATCH 3/3] test_post_batch_calls_get_conn_if_no_base_url --- .../providers/apache/livy/hooks/test_livy.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/providers/apache/livy/hooks/test_livy.py b/tests/providers/apache/livy/hooks/test_livy.py index 89c1c4c8a4b0c..e66e3bb125667 100644 --- a/tests/providers/apache/livy/hooks/test_livy.py +++ b/tests/providers/apache/livy/hooks/test_livy.py @@ -22,6 +22,7 @@ import multidict import pytest +import requests from aiohttp import ClientResponseError, RequestInfo from requests.exceptions import RequestException @@ -288,6 +289,25 @@ def test_get_batch_fail(self, requests_mock): with pytest.raises(AirflowException): LivyHook().get_batch(BATCH_ID) + @patch("airflow.providers.apache.livy.hooks.livy.LivyHook.run_method") + @patch("airflow.providers.apache.livy.hooks.livy.LivyHook.get_conn") + def test_post_batch_calls_get_conn_if_no_batch_id(self, mock_run_method, mock_get_conn): + # mock run_method to get rid of call get_conn in it + mock_response = MagicMock(resp=requests.Response) + mock_response.json.return_value = {"id": BATCH_ID, "state": BatchState.STARTING.value, "log": []} + mock_run_method.return_value = mock_response + + hook = LivyHook() + + # base_url is not set + hook.post_batch(file="sparkapp") + mock_get_conn.assert_called_once() + + # base_url is set + hook.base_url = "//livy:8998" + hook.post_batch(file="sparkapp") + mock_get_conn.not_called() + def test_invalid_uri(self): with pytest.raises(RequestException): LivyHook(livy_conn_id="invalid_uri").post_batch(file="sparkapp")