From af3c733a475b20f57c74b6cbd53adf93b0c95253 Mon Sep 17 00:00:00 2001 From: Dov Benyomin Sohacheski Date: Sun, 30 Oct 2022 14:09:02 +0200 Subject: [PATCH 1/7] Add spark's app_id to xcom output --- airflow/providers/apache/livy/operators/livy.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/airflow/providers/apache/livy/operators/livy.py b/airflow/providers/apache/livy/operators/livy.py index 763773ef04856..da9ffa6348e1e 100644 --- a/airflow/providers/apache/livy/operators/livy.py +++ b/airflow/providers/apache/livy/operators/livy.py @@ -140,6 +140,12 @@ def execute(self, context: Context) -> Any: if self._polling_interval > 0: self.poll_for_termination(self._batch_id) + + ti = context["ti"] + ti.xcom_push( + key="app_id", + value=self.get_hook().get_batch(self._batch_id)["appId"] + ) return self._batch_id From 816d1f6647bc3b8aac3e303291414db37b07cddc Mon Sep 17 00:00:00 2001 From: Dov Benyomin Sohacheski Date: Sun, 30 Oct 2022 14:38:47 +0200 Subject: [PATCH 2/7] added context mock --- tests/providers/apache/livy/operators/test_livy.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/providers/apache/livy/operators/test_livy.py b/tests/providers/apache/livy/operators/test_livy.py index 8f4e7e04f5b23..b6b82ebdf7128 100644 --- a/tests/providers/apache/livy/operators/test_livy.py +++ b/tests/providers/apache/livy/operators/test_livy.py @@ -45,6 +45,7 @@ def setUp(self): conn_id="livyunittest", conn_type="livy", host="localhost:8998", port="8998", schema="http" ) ) + self.mock_context = dict(ti=MagicMock()) @patch( "airflow.providers.apache.livy.operators.livy.LivyHook.dump_batch_logs", @@ -115,7 +116,7 @@ def test_execution(self, mock_post, mock_get, mock_dump_logs): dag=self.dag, task_id="livy_example", ) - task.execute(context={}) + task.execute(context=self.mock_context) call_args = {k: v for k, v in mock_post.call_args[1].items() if v} assert call_args == {"file": "sparkapp"} @@ -129,7 +130,7 @@ def test_execution_with_extra_options(self, mock_post): file="sparkapp", dag=self.dag, task_id="livy_example", extra_options=extra_options ) - task.execute(context={}) + task.execute(context=self.mock_context) assert task.get_hook().extra_options == extra_options @@ -139,7 +140,7 @@ def test_deletion(self, mock_post, mock_delete): task = LivyOperator( livy_conn_id="livyunittest", file="sparkapp", dag=self.dag, task_id="livy_example" ) - task.execute(context={}) + task.execute(context=self.mock_context) task.kill() mock_delete.assert_called_once_with(BATCH_ID) @@ -167,7 +168,7 @@ def test_log_dump(self, mock_post, mock_get_logs, mock_get): polling_interval=1, ) with self.assertLogs(task.get_hook().log, level=logging.INFO) as cm: - task.execute(context={}) + task.execute(context=self.mock_context) assert "INFO:airflow.providers.apache.livy.hooks.livy.LivyHook:first_line" in cm.output assert "INFO:airflow.providers.apache.livy.hooks.livy.LivyHook:second_line" in cm.output assert "INFO:airflow.providers.apache.livy.hooks.livy.LivyHook:third_line" in cm.output From 0faa59559ed7c584cb01feaef1c21bdacfbf2424 Mon Sep 17 00:00:00 2001 From: Dov Benyomin Sohacheski Date: Sun, 30 Oct 2022 18:42:20 +0200 Subject: [PATCH 3/7] added tests --- .../providers/apache/livy/operators/livy.py | 7 +++--- .../apache/livy/operators/test_livy.py | 22 ++++++++++++------- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/airflow/providers/apache/livy/operators/livy.py b/airflow/providers/apache/livy/operators/livy.py index da9ffa6348e1e..07705e61117ee 100644 --- a/airflow/providers/apache/livy/operators/livy.py +++ b/airflow/providers/apache/livy/operators/livy.py @@ -140,10 +140,9 @@ def execute(self, context: Context) -> Any: if self._polling_interval > 0: self.poll_for_termination(self._batch_id) - - ti = context["ti"] - ti.xcom_push( - key="app_id", + + context["ti"].xcom_push( + key="app_id", value=self.get_hook().get_batch(self._batch_id)["appId"] ) diff --git a/tests/providers/apache/livy/operators/test_livy.py b/tests/providers/apache/livy/operators/test_livy.py index b6b82ebdf7128..836b24e823910 100644 --- a/tests/providers/apache/livy/operators/test_livy.py +++ b/tests/providers/apache/livy/operators/test_livy.py @@ -30,9 +30,9 @@ from airflow.utils import db, timezone DEFAULT_DATE = timezone.datetime(2017, 1, 1) -mock_livy_client = MagicMock() - BATCH_ID = 100 +APP_ID = "application_1433865536131_34483" +GET_BATCH = {"appId": APP_ID} LOG_RESPONSE = {"total": 3, "log": ["first_line", "second_line", "third_line"]} @@ -53,7 +53,6 @@ def setUp(self): ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch_state") def test_poll_for_termination(self, mock_livy, mock_dump_logs): - state_list = 2 * [BatchState.RUNNING] + [BatchState.SUCCESS] def side_effect(_, retry_args): @@ -78,7 +77,6 @@ def side_effect(_, retry_args): ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch_state") def test_poll_for_termination_fail(self, mock_livy, mock_dump_logs): - state_list = 2 * [BatchState.RUNNING] + [BatchState.ERROR] def side_effect(_, retry_args): @@ -108,7 +106,8 @@ def side_effect(_, retry_args): return_value=BatchState.SUCCESS, ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch", return_value=BATCH_ID) - def test_execution(self, mock_post, mock_get, mock_dump_logs): + @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch", return_value=GET_BATCH) + def test_execution(self, mock_get_batch, mock_post, mock_get, mock_dump_logs): task = LivyOperator( livy_conn_id="livyunittest", file="sparkapp", @@ -122,9 +121,14 @@ def test_execution(self, mock_post, mock_get, mock_dump_logs): assert call_args == {"file": "sparkapp"} mock_get.assert_called_once_with(BATCH_ID, retry_args=None) mock_dump_logs.assert_called_once_with(BATCH_ID) + mock_get_batch.assert_called_once_with(BATCH_ID) + self.mock_context["ti"].xcom_push.assert_called_once_with( + key="app_id", value=APP_ID + ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch") - def test_execution_with_extra_options(self, mock_post): + @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch", return_value=GET_BATCH) + def test_execution_with_extra_options(self, mock_get_batch, mock_post): extra_options = {"check_response": True} task = LivyOperator( file="sparkapp", dag=self.dag, task_id="livy_example", extra_options=extra_options @@ -136,7 +140,8 @@ def test_execution_with_extra_options(self, mock_post): @patch("airflow.providers.apache.livy.operators.livy.LivyHook.delete_batch") @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch", return_value=BATCH_ID) - def test_deletion(self, mock_post, mock_delete): + @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch", return_value=GET_BATCH) + def test_deletion(self, mock_get_batch, mock_post, mock_delete): task = LivyOperator( livy_conn_id="livyunittest", file="sparkapp", dag=self.dag, task_id="livy_example" ) @@ -159,7 +164,8 @@ def test_injected_hook(self): ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch_logs", return_value=LOG_RESPONSE) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch", return_value=BATCH_ID) - def test_log_dump(self, mock_post, mock_get_logs, mock_get): + @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch", return_value=GET_BATCH) + def test_log_dump(self, mock_get_batch, mock_post, mock_get_logs, mock_get): task = LivyOperator( livy_conn_id="livyunittest", file="sparkapp", From 63e81a33f6804e5c83aa1c55c40d6b7a6e4479ca Mon Sep 17 00:00:00 2001 From: Dov Benyomin Sohacheski Date: Sun, 30 Oct 2022 14:09:02 +0200 Subject: [PATCH 4/7] Add spark's app_id to xcom output --- airflow/providers/apache/livy/operators/livy.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/airflow/providers/apache/livy/operators/livy.py b/airflow/providers/apache/livy/operators/livy.py index 763773ef04856..da9ffa6348e1e 100644 --- a/airflow/providers/apache/livy/operators/livy.py +++ b/airflow/providers/apache/livy/operators/livy.py @@ -140,6 +140,12 @@ def execute(self, context: Context) -> Any: if self._polling_interval > 0: self.poll_for_termination(self._batch_id) + + ti = context["ti"] + ti.xcom_push( + key="app_id", + value=self.get_hook().get_batch(self._batch_id)["appId"] + ) return self._batch_id From b0020a6efe49adeb3eca0ffdbf79527e4dfe8aa7 Mon Sep 17 00:00:00 2001 From: Dov Benyomin Sohacheski Date: Sun, 30 Oct 2022 14:38:47 +0200 Subject: [PATCH 5/7] added context mock --- tests/providers/apache/livy/operators/test_livy.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/providers/apache/livy/operators/test_livy.py b/tests/providers/apache/livy/operators/test_livy.py index 8f4e7e04f5b23..b6b82ebdf7128 100644 --- a/tests/providers/apache/livy/operators/test_livy.py +++ b/tests/providers/apache/livy/operators/test_livy.py @@ -45,6 +45,7 @@ def setUp(self): conn_id="livyunittest", conn_type="livy", host="localhost:8998", port="8998", schema="http" ) ) + self.mock_context = dict(ti=MagicMock()) @patch( "airflow.providers.apache.livy.operators.livy.LivyHook.dump_batch_logs", @@ -115,7 +116,7 @@ def test_execution(self, mock_post, mock_get, mock_dump_logs): dag=self.dag, task_id="livy_example", ) - task.execute(context={}) + task.execute(context=self.mock_context) call_args = {k: v for k, v in mock_post.call_args[1].items() if v} assert call_args == {"file": "sparkapp"} @@ -129,7 +130,7 @@ def test_execution_with_extra_options(self, mock_post): file="sparkapp", dag=self.dag, task_id="livy_example", extra_options=extra_options ) - task.execute(context={}) + task.execute(context=self.mock_context) assert task.get_hook().extra_options == extra_options @@ -139,7 +140,7 @@ def test_deletion(self, mock_post, mock_delete): task = LivyOperator( livy_conn_id="livyunittest", file="sparkapp", dag=self.dag, task_id="livy_example" ) - task.execute(context={}) + task.execute(context=self.mock_context) task.kill() mock_delete.assert_called_once_with(BATCH_ID) @@ -167,7 +168,7 @@ def test_log_dump(self, mock_post, mock_get_logs, mock_get): polling_interval=1, ) with self.assertLogs(task.get_hook().log, level=logging.INFO) as cm: - task.execute(context={}) + task.execute(context=self.mock_context) assert "INFO:airflow.providers.apache.livy.hooks.livy.LivyHook:first_line" in cm.output assert "INFO:airflow.providers.apache.livy.hooks.livy.LivyHook:second_line" in cm.output assert "INFO:airflow.providers.apache.livy.hooks.livy.LivyHook:third_line" in cm.output From c1b750621ef8a4b3af0d11c724fc782d353d29c5 Mon Sep 17 00:00:00 2001 From: Dov Benyomin Sohacheski Date: Sun, 30 Oct 2022 18:42:20 +0200 Subject: [PATCH 6/7] added tests --- .../providers/apache/livy/operators/livy.py | 7 +++--- .../apache/livy/operators/test_livy.py | 22 ++++++++++++------- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/airflow/providers/apache/livy/operators/livy.py b/airflow/providers/apache/livy/operators/livy.py index da9ffa6348e1e..07705e61117ee 100644 --- a/airflow/providers/apache/livy/operators/livy.py +++ b/airflow/providers/apache/livy/operators/livy.py @@ -140,10 +140,9 @@ def execute(self, context: Context) -> Any: if self._polling_interval > 0: self.poll_for_termination(self._batch_id) - - ti = context["ti"] - ti.xcom_push( - key="app_id", + + context["ti"].xcom_push( + key="app_id", value=self.get_hook().get_batch(self._batch_id)["appId"] ) diff --git a/tests/providers/apache/livy/operators/test_livy.py b/tests/providers/apache/livy/operators/test_livy.py index b6b82ebdf7128..836b24e823910 100644 --- a/tests/providers/apache/livy/operators/test_livy.py +++ b/tests/providers/apache/livy/operators/test_livy.py @@ -30,9 +30,9 @@ from airflow.utils import db, timezone DEFAULT_DATE = timezone.datetime(2017, 1, 1) -mock_livy_client = MagicMock() - BATCH_ID = 100 +APP_ID = "application_1433865536131_34483" +GET_BATCH = {"appId": APP_ID} LOG_RESPONSE = {"total": 3, "log": ["first_line", "second_line", "third_line"]} @@ -53,7 +53,6 @@ def setUp(self): ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch_state") def test_poll_for_termination(self, mock_livy, mock_dump_logs): - state_list = 2 * [BatchState.RUNNING] + [BatchState.SUCCESS] def side_effect(_, retry_args): @@ -78,7 +77,6 @@ def side_effect(_, retry_args): ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch_state") def test_poll_for_termination_fail(self, mock_livy, mock_dump_logs): - state_list = 2 * [BatchState.RUNNING] + [BatchState.ERROR] def side_effect(_, retry_args): @@ -108,7 +106,8 @@ def side_effect(_, retry_args): return_value=BatchState.SUCCESS, ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch", return_value=BATCH_ID) - def test_execution(self, mock_post, mock_get, mock_dump_logs): + @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch", return_value=GET_BATCH) + def test_execution(self, mock_get_batch, mock_post, mock_get, mock_dump_logs): task = LivyOperator( livy_conn_id="livyunittest", file="sparkapp", @@ -122,9 +121,14 @@ def test_execution(self, mock_post, mock_get, mock_dump_logs): assert call_args == {"file": "sparkapp"} mock_get.assert_called_once_with(BATCH_ID, retry_args=None) mock_dump_logs.assert_called_once_with(BATCH_ID) + mock_get_batch.assert_called_once_with(BATCH_ID) + self.mock_context["ti"].xcom_push.assert_called_once_with( + key="app_id", value=APP_ID + ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch") - def test_execution_with_extra_options(self, mock_post): + @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch", return_value=GET_BATCH) + def test_execution_with_extra_options(self, mock_get_batch, mock_post): extra_options = {"check_response": True} task = LivyOperator( file="sparkapp", dag=self.dag, task_id="livy_example", extra_options=extra_options @@ -136,7 +140,8 @@ def test_execution_with_extra_options(self, mock_post): @patch("airflow.providers.apache.livy.operators.livy.LivyHook.delete_batch") @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch", return_value=BATCH_ID) - def test_deletion(self, mock_post, mock_delete): + @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch", return_value=GET_BATCH) + def test_deletion(self, mock_get_batch, mock_post, mock_delete): task = LivyOperator( livy_conn_id="livyunittest", file="sparkapp", dag=self.dag, task_id="livy_example" ) @@ -159,7 +164,8 @@ def test_injected_hook(self): ) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch_logs", return_value=LOG_RESPONSE) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch", return_value=BATCH_ID) - def test_log_dump(self, mock_post, mock_get_logs, mock_get): + @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch", return_value=GET_BATCH) + def test_log_dump(self, mock_get_batch, mock_post, mock_get_logs, mock_get): task = LivyOperator( livy_conn_id="livyunittest", file="sparkapp", From 6f7148b2bba09a50359ef29ec06d8ecf3dd001f2 Mon Sep 17 00:00:00 2001 From: Dov Benyomin Sohacheski Date: Sun, 30 Oct 2022 21:50:42 +0200 Subject: [PATCH 7/7] fixed static styling --- airflow/providers/apache/livy/operators/livy.py | 5 +---- tests/providers/apache/livy/operators/test_livy.py | 4 +--- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/airflow/providers/apache/livy/operators/livy.py b/airflow/providers/apache/livy/operators/livy.py index 07705e61117ee..5d2293b42d280 100644 --- a/airflow/providers/apache/livy/operators/livy.py +++ b/airflow/providers/apache/livy/operators/livy.py @@ -141,10 +141,7 @@ def execute(self, context: Context) -> Any: if self._polling_interval > 0: self.poll_for_termination(self._batch_id) - context["ti"].xcom_push( - key="app_id", - value=self.get_hook().get_batch(self._batch_id)["appId"] - ) + context["ti"].xcom_push(key="app_id", value=self.get_hook().get_batch(self._batch_id)["appId"]) return self._batch_id diff --git a/tests/providers/apache/livy/operators/test_livy.py b/tests/providers/apache/livy/operators/test_livy.py index 836b24e823910..d49cba1c9cf41 100644 --- a/tests/providers/apache/livy/operators/test_livy.py +++ b/tests/providers/apache/livy/operators/test_livy.py @@ -122,9 +122,7 @@ def test_execution(self, mock_get_batch, mock_post, mock_get, mock_dump_logs): mock_get.assert_called_once_with(BATCH_ID, retry_args=None) mock_dump_logs.assert_called_once_with(BATCH_ID) mock_get_batch.assert_called_once_with(BATCH_ID) - self.mock_context["ti"].xcom_push.assert_called_once_with( - key="app_id", value=APP_ID - ) + self.mock_context["ti"].xcom_push.assert_called_once_with(key="app_id", value=APP_ID) @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch") @patch("airflow.providers.apache.livy.operators.livy.LivyHook.get_batch", return_value=GET_BATCH)