Resolve deprecations in the tests for Google Dataflow operators (#4…

…0327)
apache · Jun 20, 2024 · cfb83c4 · cfb83c4
1 parent dd69748
commit cfb83c4
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 40 deletions.
diff --git a/tests/deprecations_ignore.yml b/tests/deprecations_ignore.yml
@@ -182,15 +182,6 @@
 - tests/providers/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithoutDefaultProjectIdHook::test_get_pipeline_job
 - tests/providers/google/cloud/hooks/vertex_ai/test_custom_job.py::TestCustomJobWithoutDefaultProjectIdHook::test_list_pipeline_jobs
 - tests/providers/google/cloud/operators/test_automl.py::TestAutoMLTrainModelOperator::test_execute
-- tests/providers/google/cloud/operators/test_dataflow.py::TestDataflowCreateJavaJobOperator::test_check_job_not_running_exec
-- tests/providers/google/cloud/operators/test_dataflow.py::TestDataflowCreateJavaJobOperator::test_check_job_running_exec
-- tests/providers/google/cloud/operators/test_dataflow.py::TestDataflowCreateJavaJobOperator::test_check_multiple_job_exec
-- tests/providers/google/cloud/operators/test_dataflow.py::TestDataflowCreateJavaJobOperator::test_exec
-- tests/providers/google/cloud/operators/test_dataflow.py::TestDataflowCreateJavaJobOperator::test_init
-- tests/providers/google/cloud/operators/test_dataflow.py::TestDataflowCreateJavaJobOperatorWithLocal::test_check_job_not_running_exec
-- tests/providers/google/cloud/operators/test_dataflow.py::TestDataflowCreateJavaJobOperatorWithLocal::test_init
-- tests/providers/google/cloud/operators/test_dataflow.py::TestDataflowCreatePythonJobOperator::test_exec
-- tests/providers/google/cloud/operators/test_dataflow.py::TestDataflowCreatePythonJobOperator::test_init
 - tests/providers/google/cloud/operators/test_datapipeline.py::TestCreateDataPipelineOperator::test_execute
 - tests/providers/google/cloud/operators/test_datapipeline.py::TestRunDataPipelineOperator::test_execute
 - tests/providers/google/cloud/operators/test_dataproc.py::TestDataProcHadoopOperator::test_execute

diff --git a/tests/providers/google/cloud/operators/test_dataflow.py b/tests/providers/google/cloud/operators/test_dataflow.py
@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import copy
+import warnings
 from copy import deepcopy
 from unittest import mock
 
@@ -26,7 +27,7 @@
 from googleapiclient.errors import HttpError
 
 import airflow
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.dataflow import (
     DEFAULT_DATAFLOW_LOCATION,
     DataflowJobStatus,
@@ -134,16 +135,18 @@
 
 class TestDataflowCreatePythonJobOperator:
     def setup_method(self):
-        self.dataflow = DataflowCreatePythonJobOperator(
-            task_id=TASK_ID,
-            py_file=PY_FILE,
-            job_name=JOB_NAME,
-            py_options=PY_OPTIONS,
-            dataflow_default_options=DEFAULT_OPTIONS_PYTHON,
-            options=ADDITIONAL_OPTIONS,
-            poll_sleep=POLL_SLEEP,
-            location=TEST_LOCATION,
-        )
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", AirflowProviderDeprecationWarning)
+            self.dataflow = DataflowCreatePythonJobOperator(
+                task_id=TASK_ID,
+                py_file=PY_FILE,
+                job_name=JOB_NAME,
+                py_options=PY_OPTIONS,
+                dataflow_default_options=DEFAULT_OPTIONS_PYTHON,
+                options=ADDITIONAL_OPTIONS,
+                poll_sleep=POLL_SLEEP,
+                location=TEST_LOCATION,
+            )
         self.expected_airflow_version = "v" + airflow.version.version.replace(".", "-").replace("+", "-")
 
     def test_init(self):
@@ -214,16 +217,18 @@ def test_exec(self, gcs_hook, dataflow_hook_mock, beam_hook_mock, mock_callback_
 
 class TestDataflowCreateJavaJobOperator:
     def setup_method(self):
-        self.dataflow = DataflowCreateJavaJobOperator(
-            task_id=TASK_ID,
-            jar=JAR_FILE,
-            job_name=JOB_NAME,
-            job_class=JOB_CLASS,
-            dataflow_default_options=DEFAULT_OPTIONS_JAVA,
-            options=ADDITIONAL_OPTIONS,
-            poll_sleep=POLL_SLEEP,
-            location=TEST_LOCATION,
-        )
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", AirflowProviderDeprecationWarning)
+            self.dataflow = DataflowCreateJavaJobOperator(
+                task_id=TASK_ID,
+                jar=JAR_FILE,
+                job_name=JOB_NAME,
+                job_class=JOB_CLASS,
+                dataflow_default_options=DEFAULT_OPTIONS_JAVA,
+                options=ADDITIONAL_OPTIONS,
+                poll_sleep=POLL_SLEEP,
+                location=TEST_LOCATION,
+            )
         self.expected_airflow_version = "v" + airflow.version.version.replace(".", "-").replace("+", "-")
 
     def test_init(self):
@@ -421,16 +426,18 @@ def set_is_job_dataflow_running_variables(*args, **kwargs):
 
 class TestDataflowCreateJavaJobOperatorWithLocal:
     def setup_method(self):
-        self.dataflow = DataflowCreateJavaJobOperator(
-            task_id=TASK_ID,
-            jar=LOCAL_JAR_FILE,
-            job_name=JOB_NAME,
-            job_class=JOB_CLASS,
-            dataflow_default_options=DEFAULT_OPTIONS_JAVA,
-            options=ADDITIONAL_OPTIONS,
-            poll_sleep=POLL_SLEEP,
-            location=TEST_LOCATION,
-        )
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", AirflowProviderDeprecationWarning)
+            self.dataflow = DataflowCreateJavaJobOperator(
+                task_id=TASK_ID,
+                jar=LOCAL_JAR_FILE,
+                job_name=JOB_NAME,
+                job_class=JOB_CLASS,
+                dataflow_default_options=DEFAULT_OPTIONS_JAVA,
+                options=ADDITIONAL_OPTIONS,
+                poll_sleep=POLL_SLEEP,
+                location=TEST_LOCATION,
+            )
         self.expected_airflow_version = "v" + airflow.version.version.replace(".", "-").replace("+", "-")
 
     def test_init(self):