Resolve deprecations in SubDagOperator (apache#40179)

okirialbert · Jun 12, 2024 · 794678f · 794678f
1 parent 1372e10
commit 794678f
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 26 deletions.
diff --git a/tests/deprecations_ignore.yml b/tests/deprecations_ignore.yml
@@ -47,16 +47,6 @@
 
 # Operators
 - tests/operators/test_latest_only_operator.py::TestLatestOnlyOperator::test_skipping_non_latest
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_execute_create_dagrun_wait_until_success
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_execute_create_dagrun_with_conf
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_execute_dagrun_failed
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_execute_skip_if_dagrun_success
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_rerun_failed_subdag
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_subdag_in_context_manager
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_subdag_name
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_subdag_pools
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_subdag_pools_no_possible_conflict
-- tests/operators/test_subdag_operator.py::TestSubDagOperator::test_subdag_with_propagate_skipped_state
 
 
 # Serialization

diff --git a/tests/operators/test_subdag_operator.py b/tests/operators/test_subdag_operator.py
@@ -23,7 +23,7 @@
 import pytest
 
 import airflow
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, RemovedInAirflow3Warning
 from airflow.models.dag import DAG
 from airflow.models.dagrun import DagRun
 from airflow.models.taskinstance import TaskInstance
@@ -41,6 +41,8 @@
 
 default_args = {"start_date": DEFAULT_DATE}
 
+WARNING_MESSAGE = """This class is deprecated. Please use `airflow.utils.task_group.TaskGroup`."""
+
 
 class TestSubDagOperator:
     def setup_method(self):
@@ -65,7 +67,8 @@ def test_subdag_name(self):
         subdag_bad2 = DAG("bad.test", default_args=default_args)
         subdag_bad3 = DAG("bad.bad", default_args=default_args)
 
-        SubDagOperator(task_id="test", dag=dag, subdag=subdag_good)
+        with pytest.warns(RemovedInAirflow3Warning, match=WARNING_MESSAGE):
+            SubDagOperator(task_id="test", dag=dag, subdag=subdag_good)
         with pytest.raises(AirflowException):
             SubDagOperator(task_id="test", dag=dag, subdag=subdag_bad1)
         with pytest.raises(AirflowException):
@@ -79,7 +82,8 @@ def test_subdag_in_context_manager(self):
         """
         with DAG("parent", default_args=default_args) as dag:
             subdag = DAG("parent.test", default_args=default_args)
-            op = SubDagOperator(task_id="test", subdag=subdag)
+            with pytest.warns(RemovedInAirflow3Warning, match=WARNING_MESSAGE):
+                op = SubDagOperator(task_id="test", subdag=subdag)
 
             assert op.dag == dag
             assert op.subdag == subdag
@@ -105,7 +109,8 @@ def test_subdag_pools(self):
 
         # recreate dag because failed subdagoperator was already added
         dag = DAG("parent", default_args=default_args)
-        SubDagOperator(task_id="child", dag=dag, subdag=subdag, pool="test_pool_10")
+        with pytest.warns(RemovedInAirflow3Warning, match=WARNING_MESSAGE):
+            SubDagOperator(task_id="child", dag=dag, subdag=subdag, pool="test_pool_10")
 
         session.delete(pool_1)
         session.delete(pool_10)
@@ -129,7 +134,8 @@ def test_subdag_pools_no_possible_conflict(self):
         EmptyOperator(task_id="dummy", dag=subdag, pool="test_pool_10")
 
         mock_session = Mock()
-        SubDagOperator(task_id="child", dag=dag, subdag=subdag, pool="test_pool_1", session=mock_session)
+        with pytest.warns(RemovedInAirflow3Warning, match=WARNING_MESSAGE):
+            SubDagOperator(task_id="child", dag=dag, subdag=subdag, pool="test_pool_1", session=mock_session)
         assert not mock_session.query.called
 
         session.delete(pool_1)
@@ -143,7 +149,8 @@ def test_execute_create_dagrun_wait_until_success(self):
         """
         dag = DAG("parent", default_args=default_args)
         subdag = DAG("parent.test", default_args=default_args)
-        subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1)
+        with pytest.warns(RemovedInAirflow3Warning, match=WARNING_MESSAGE):
+            subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1)
 
         subdag.create_dagrun = Mock()
         subdag.create_dagrun.return_value = self.dag_run_running
@@ -180,7 +187,8 @@ def test_execute_create_dagrun_with_conf(self):
         conf = {"key": "value"}
         dag = DAG("parent", default_args=default_args)
         subdag = DAG("parent.test", default_args=default_args)
-        subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1, conf=conf)
+        with pytest.warns(RemovedInAirflow3Warning, match=WARNING_MESSAGE):
+            subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1, conf=conf)
 
         subdag.create_dagrun = Mock()
         subdag.create_dagrun.return_value = self.dag_run_running
@@ -215,7 +223,8 @@ def test_execute_dagrun_failed(self):
         """
         dag = DAG("parent", default_args=default_args)
         subdag = DAG("parent.test", default_args=default_args)
-        subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1)
+        with pytest.warns(RemovedInAirflow3Warning, match=WARNING_MESSAGE):
+            subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1)
 
         subdag.create_dagrun = Mock()
         subdag.create_dagrun.return_value = self.dag_run_running
@@ -242,7 +251,8 @@ def test_execute_skip_if_dagrun_success(self):
         subdag = DAG("parent.test", default_args=default_args)
 
         subdag.create_dagrun = Mock()
-        subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1)
+        with pytest.warns(RemovedInAirflow3Warning, match=WARNING_MESSAGE):
+            subdag_task = SubDagOperator(task_id="test", subdag=subdag, dag=dag, poke_interval=1)
         subdag_task._get_dagrun = Mock()
         subdag_task._get_dagrun.return_value = self.dag_run_success
 
@@ -278,7 +288,9 @@ def test_rerun_failed_subdag(self, dag_maker):
             dummy_task_instance.refresh_from_task(dummy_task)
             dummy_task_instance.state == State.FAILED
 
-            with dag_maker("parent", default_args=default_args, session=session):
+            with dag_maker("parent", default_args=default_args, session=session), pytest.warns(
+                RemovedInAirflow3Warning, match=WARNING_MESSAGE
+            ):
                 subdag_task = SubDagOperator(task_id="test", subdag=subdag, poke_interval=1)
             dag_maker.create_dagrun(execution_date=DEFAULT_DATE, run_type=DagRunType.SCHEDULED)
 
@@ -320,12 +332,13 @@ def test_subdag_with_propagate_skipped_state(
         dag_maker.create_dagrun(execution_date=DEFAULT_DATE)
 
         with dag_maker("parent", default_args=default_args):
-            subdag_task = SubDagOperator(
-                task_id="test",
-                subdag=subdag,
-                poke_interval=1,
-                propagate_skipped_state=propagate_option,
-            )
+            with pytest.warns(RemovedInAirflow3Warning, match=WARNING_MESSAGE):
+                subdag_task = SubDagOperator(
+                    task_id="test",
+                    subdag=subdag,
+                    poke_interval=1,
+                    propagate_skipped_state=propagate_option,
+                )
             dummy_dag_task = EmptyOperator(task_id="dummy_dag")
             subdag_task >> dummy_dag_task
         dag_run = dag_maker.create_dagrun(execution_date=DEFAULT_DATE)