Resolve deprecations in apache.livy provider tests (#39516)

apache · May 9, 2024 · 9fa2d6d · 9fa2d6d
1 parent 5c6d379
commit 9fa2d6d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/tests/deprecations_ignore.yml b/tests/deprecations_ignore.yml
@@ -332,7 +332,6 @@
 - tests/providers/amazon/aws/triggers/test_redshift_cluster.py::TestRedshiftClusterTrigger::test_redshift_cluster_sensor_trigger_resuming_status
 - tests/providers/amazon/aws/triggers/test_redshift_cluster.py::TestRedshiftClusterTrigger::test_redshift_cluster_sensor_trigger_success
 - tests/providers/amazon/aws/utils/test_connection_wrapper.py::TestAwsConnectionWrapper::test_get_endpoint_url_from_extra
-- tests/providers/apache/livy/operators/test_livy.py::TestLivyOperator::test_execution_with_extra_options
 - tests/providers/apache/spark/operators/test_spark_sql.py::TestSparkSqlOperator::test_execute
 - tests/providers/atlassian/jira/operators/test_jira.py::TestJiraOperator::test_issue_search
 - tests/providers/atlassian/jira/operators/test_jira.py::TestJiraOperator::test_project_issue_count

diff --git a/tests/providers/apache/livy/operators/test_livy.py b/tests/providers/apache/livy/operators/test_livy.py
@@ -21,7 +21,7 @@
 
 import pytest
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models import Connection
 from airflow.models.dag import DAG
 from airflow.providers.apache.livy.hooks.livy import BatchState
@@ -134,7 +134,7 @@ def test_execution_with_extra_options(self, mock_get_batch, mock_post):
 
         task.execute(context=self.mock_context)
 
-        assert task.get_hook().extra_options == extra_options
+        assert task.hook.extra_options == extra_options
 
     @patch("airflow.providers.apache.livy.operators.livy.LivyHook.delete_batch")
     @patch("airflow.providers.apache.livy.operators.livy.LivyHook.post_batch", return_value=BATCH_ID)
@@ -417,6 +417,12 @@ def test_execute_complete_timeout(self, mock_delete, mock_post):
         mock_delete.assert_called_once_with(BATCH_ID)
         self.mock_context["ti"].xcom_push.assert_not_called()
 
+    def test_deprecated_get_hook(self):
+        op = LivyOperator(task_id="livy_example", file="sparkapp")
+        with pytest.warns(AirflowProviderDeprecationWarning, match="use `hook` property instead"):
+            hook = op.get_hook()
+        assert hook is op.hook
+
 
 @pytest.mark.db_test
 def test_spark_params_templating(create_task_instance_of_operator):