bugfix/livy-set-base-url

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)

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")