diff --git a/airflow/providers/apache/livy/hooks/livy.py b/airflow/providers/apache/livy/hooks/livy.py index 53a86b5cfb198..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 is 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) 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")