From cfb83c4b160f498e33e249fb92692d74cf0adcdc Mon Sep 17 00:00:00 2001 From: Bora Berke Sahin <67373739+boraberke@users.noreply.github.com> Date: Thu, 20 Jun 2024 13:18:37 +0300 Subject: [PATCH] Resolve deprecations in the tests for `Google Dataflow` operators (#40327) --- tests/deprecations_ignore.yml | 9 --- .../google/cloud/operators/test_dataflow.py | 69 ++++++++++--------- 2 files changed, 38 insertions(+), 40 deletions(-) diff --git a/tests/deprecations_ignore.yml b/tests/deprecations_ignore.yml index 1806a5c7164b..172fc78418f4 100644 --- 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 index 7b8aeaf3fd5e..a024ffd7a750 100644 --- 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):