jedcunningham commented on code in PR #30350: URL: https://github.com/apache/airflow/pull/30350#discussion_r1153814185
########## airflow/utils/setup_teardown.py: ########## @@ -0,0 +1,132 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from airflow.models.operator import Operator + + +class SetupTeardownContext: + """Context manager for setup/teardown XComArg.""" Review Comment: ```suggestion """Context manager for setup/teardown tasks.""" ``` or something ########## tests/decorators/test_setup_teardown.py: ########## @@ -211,3 +210,354 @@ def mytask2(): assert dag.task_group.children["mytask"]._is_teardown is False assert dag.task_group.children["mytask2"]._is_setup is False assert dag.task_group.children["mytask2"]._is_teardown is False + + def test_setup_teardown_as_context_manager_normal_tasks_rel_set_downstream(self, dag_maker): + """ + Test that setup >> teardown tasks are set up correctly when used as context managers + and the normal tasks are set up with >> relations. + """ + + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @teardown + def teardowntask(): + print("teardown") + + @teardown + def teardowntask2(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask() >> teardowntask(): + with setuptask2() >> teardowntask2(): + mytask() >> mytask2() + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"teardowntask", "setuptask2"} + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"teardowntask2", "mytask"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask2"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"teardowntask2"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"setuptask", "teardowntask2"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + assert dag.task_group.children["teardowntask2"].upstream_task_ids == {"setuptask2", "mytask2"} + assert dag.task_group.children["teardowntask2"].downstream_task_ids == {"teardowntask"} + + def test_setup_teardown_as_context_manager_normal_tasks_rel_set_upstream(self, dag_maker): + """ + Test that setup >> teardown tasks are set up correctly when used as context managers + and the normal tasks are set up with << relations. + """ + + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @teardown + def teardowntask(): + print("teardown") + + @teardown + def teardowntask2(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask() >> teardowntask(): + with setuptask2() >> teardowntask2(): + mytask() << mytask2() + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"teardowntask", "setuptask2"} + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"teardowntask2", "mytask2"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"teardowntask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"setuptask2"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"setuptask", "teardowntask2"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + assert dag.task_group.children["teardowntask2"].upstream_task_ids == {"setuptask2", "mytask"} + assert dag.task_group.children["teardowntask2"].downstream_task_ids == {"teardowntask"} + + def test_normal_task_raises_when_used_as_context_managers(self, dag_maker): + @task() + def mytask(): + print("mytask") + + with dag_maker(): + with pytest.raises( + AirflowException, match="Only setup/teardown tasks can be used as context managers." + ): + with mytask(): + pass + + def test_only_setup(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask(): + mytask() >> mytask2() + + assert len(dag.task_group.children) == 3 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert not dag.task_group.children["mytask2"].downstream_task_ids + + def test_only_teardown(self, dag_maker): + @teardown + def teardowntask(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker("foo") as dag: + with teardowntask(): + mytask() >> mytask2() + + assert len(dag.task_group.children) == 3 + assert not dag.task_group.children["mytask"].upstream_task_ids + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"teardowntask"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"mytask2"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + + def test_nested_only_setup(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @teardown + def teardowntask(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask() >> teardowntask(): + with setuptask2(): + mytask() << mytask2() + + assert len(dag.task_group.children) == 5 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"teardowntask", "setuptask2"} + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"teardowntask"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"setuptask2"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"setuptask", "mytask"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + + def test_task_in_different_setup_context(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + @task() + def mytask3(): + print("mytask") + + @task + def mytask4(): + print("mytask") + + with dag_maker() as dag: + with setuptask(): + t1 = mytask() + t2 = mytask2() + t1 >> t2 + with setuptask2(): + t3 = mytask3() + t4 = mytask4() + t2 >> t3 >> t4 + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert not dag.task_group.children["setuptask2"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"mytask3"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} Review Comment: ```suggestion assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} assert not dag.task_group.children["mytask2"].downstream_task_ids ``` ########## tests/decorators/test_setup_teardown.py: ########## @@ -211,3 +210,354 @@ def mytask2(): assert dag.task_group.children["mytask"]._is_teardown is False assert dag.task_group.children["mytask2"]._is_setup is False assert dag.task_group.children["mytask2"]._is_teardown is False + + def test_setup_teardown_as_context_manager_normal_tasks_rel_set_downstream(self, dag_maker): + """ + Test that setup >> teardown tasks are set up correctly when used as context managers + and the normal tasks are set up with >> relations. + """ + + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @teardown + def teardowntask(): + print("teardown") + + @teardown + def teardowntask2(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask() >> teardowntask(): + with setuptask2() >> teardowntask2(): + mytask() >> mytask2() + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"teardowntask", "setuptask2"} + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"teardowntask2", "mytask"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask2"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"teardowntask2"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"setuptask", "teardowntask2"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + assert dag.task_group.children["teardowntask2"].upstream_task_ids == {"setuptask2", "mytask2"} + assert dag.task_group.children["teardowntask2"].downstream_task_ids == {"teardowntask"} + + def test_setup_teardown_as_context_manager_normal_tasks_rel_set_upstream(self, dag_maker): + """ + Test that setup >> teardown tasks are set up correctly when used as context managers + and the normal tasks are set up with << relations. + """ + + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @teardown + def teardowntask(): + print("teardown") + + @teardown + def teardowntask2(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask() >> teardowntask(): + with setuptask2() >> teardowntask2(): + mytask() << mytask2() + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"teardowntask", "setuptask2"} + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"teardowntask2", "mytask2"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"teardowntask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"setuptask2"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"setuptask", "teardowntask2"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + assert dag.task_group.children["teardowntask2"].upstream_task_ids == {"setuptask2", "mytask"} + assert dag.task_group.children["teardowntask2"].downstream_task_ids == {"teardowntask"} + + def test_normal_task_raises_when_used_as_context_managers(self, dag_maker): + @task() + def mytask(): + print("mytask") + + with dag_maker(): + with pytest.raises( + AirflowException, match="Only setup/teardown tasks can be used as context managers." + ): + with mytask(): + pass + + def test_only_setup(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask(): + mytask() >> mytask2() + + assert len(dag.task_group.children) == 3 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert not dag.task_group.children["mytask2"].downstream_task_ids + + def test_only_teardown(self, dag_maker): + @teardown + def teardowntask(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker("foo") as dag: + with teardowntask(): + mytask() >> mytask2() + + assert len(dag.task_group.children) == 3 + assert not dag.task_group.children["mytask"].upstream_task_ids + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"teardowntask"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"mytask2"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + + def test_nested_only_setup(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @teardown + def teardowntask(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask() >> teardowntask(): + with setuptask2(): + mytask() << mytask2() + + assert len(dag.task_group.children) == 5 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"teardowntask", "setuptask2"} + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"teardowntask"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"setuptask2"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"setuptask", "mytask"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + + def test_task_in_different_setup_context(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + @task() + def mytask3(): + print("mytask") + + @task + def mytask4(): + print("mytask") + + with dag_maker() as dag: + with setuptask(): + t1 = mytask() + t2 = mytask2() + t1 >> t2 + with setuptask2(): + t3 = mytask3() + t4 = mytask4() + t2 >> t3 >> t4 + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert not dag.task_group.children["setuptask2"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"mytask3"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask3"].upstream_task_ids == {"mytask2", "setuptask2"} + assert dag.task_group.children["mytask3"].downstream_task_ids == {"mytask4"} + assert dag.task_group.children["mytask4"].upstream_task_ids == {"mytask3"} Review Comment: ```suggestion assert dag.task_group.children["mytask4"].upstream_task_ids == {"mytask3"} assert not dag.task_group.children["mytask4"].downstream_task_ids ``` ########## tests/decorators/test_setup_teardown.py: ########## @@ -211,3 +210,354 @@ def mytask2(): assert dag.task_group.children["mytask"]._is_teardown is False assert dag.task_group.children["mytask2"]._is_setup is False assert dag.task_group.children["mytask2"]._is_teardown is False + + def test_setup_teardown_as_context_manager_normal_tasks_rel_set_downstream(self, dag_maker): + """ + Test that setup >> teardown tasks are set up correctly when used as context managers + and the normal tasks are set up with >> relations. + """ + + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @teardown + def teardowntask(): + print("teardown") + + @teardown + def teardowntask2(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask() >> teardowntask(): + with setuptask2() >> teardowntask2(): + mytask() >> mytask2() + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"teardowntask", "setuptask2"} + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"teardowntask2", "mytask"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask2"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"teardowntask2"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"setuptask", "teardowntask2"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + assert dag.task_group.children["teardowntask2"].upstream_task_ids == {"setuptask2", "mytask2"} + assert dag.task_group.children["teardowntask2"].downstream_task_ids == {"teardowntask"} + + def test_setup_teardown_as_context_manager_normal_tasks_rel_set_upstream(self, dag_maker): + """ + Test that setup >> teardown tasks are set up correctly when used as context managers + and the normal tasks are set up with << relations. + """ + + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @teardown + def teardowntask(): + print("teardown") + + @teardown + def teardowntask2(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask() >> teardowntask(): + with setuptask2() >> teardowntask2(): + mytask() << mytask2() + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"teardowntask", "setuptask2"} + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"teardowntask2", "mytask2"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"teardowntask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"setuptask2"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"setuptask", "teardowntask2"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + assert dag.task_group.children["teardowntask2"].upstream_task_ids == {"setuptask2", "mytask"} + assert dag.task_group.children["teardowntask2"].downstream_task_ids == {"teardowntask"} + + def test_normal_task_raises_when_used_as_context_managers(self, dag_maker): + @task() + def mytask(): + print("mytask") + + with dag_maker(): + with pytest.raises( + AirflowException, match="Only setup/teardown tasks can be used as context managers." + ): + with mytask(): + pass + + def test_only_setup(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask(): + mytask() >> mytask2() + + assert len(dag.task_group.children) == 3 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert not dag.task_group.children["mytask2"].downstream_task_ids + + def test_only_teardown(self, dag_maker): + @teardown + def teardowntask(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker("foo") as dag: + with teardowntask(): + mytask() >> mytask2() + + assert len(dag.task_group.children) == 3 + assert not dag.task_group.children["mytask"].upstream_task_ids + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"teardowntask"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"mytask2"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + + def test_nested_only_setup(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @teardown + def teardowntask(): + print("teardown") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + with dag_maker() as dag: + with setuptask() >> teardowntask(): + with setuptask2(): + mytask() << mytask2() + + assert len(dag.task_group.children) == 5 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"teardowntask", "setuptask2"} + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"teardowntask"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"setuptask2"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["teardowntask"].upstream_task_ids == {"setuptask", "mytask"} + assert not dag.task_group.children["teardowntask"].downstream_task_ids + + def test_task_in_different_setup_context(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + @task() + def mytask3(): + print("mytask") + + @task + def mytask4(): + print("mytask") + + with dag_maker() as dag: + with setuptask(): + t1 = mytask() + t2 = mytask2() + t1 >> t2 + with setuptask2(): + t3 = mytask3() + t4 = mytask4() + t2 >> t3 >> t4 + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert not dag.task_group.children["setuptask2"].upstream_task_ids + assert dag.task_group.children["setuptask"].downstream_task_ids == {"mytask"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"mytask3"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask3"].upstream_task_ids == {"mytask2", "setuptask2"} + assert dag.task_group.children["mytask3"].downstream_task_ids == {"mytask4"} + assert dag.task_group.children["mytask4"].upstream_task_ids == {"mytask3"} + + def test_task_in_different_setup_context_2(self, dag_maker): + @setup + def setuptask(): + print("setup") + + @setup + def setuptask2(): + print("setup") + + @task() + def mytask(): + print("mytask") + + @task() + def mytask2(): + print("mytask") + + @task() + def mytask3(): + print("mytask") + + @task + def mytask4(): + print("mytask") + + with dag_maker() as dag: + with setuptask(): + t1 = mytask() + t2 = mytask2() + t1 >> t2 + with setuptask2(): + t3 = mytask3() + t4 = mytask4() + t2 >> t3 >> t4 + + assert len(dag.task_group.children) == 6 + assert not dag.task_group.children["setuptask"].upstream_task_ids + assert dag.task_group.children["setuptask2"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["setuptask"].downstream_task_ids == {"mytask", "setuptask2"} + assert dag.task_group.children["setuptask2"].downstream_task_ids == {"mytask3"} + assert dag.task_group.children["mytask"].upstream_task_ids == {"setuptask"} + assert dag.task_group.children["mytask"].downstream_task_ids == {"mytask2"} + assert dag.task_group.children["mytask2"].upstream_task_ids == {"mytask"} + assert dag.task_group.children["mytask2"].downstream_task_ids == {"mytask3"} + assert dag.task_group.children["mytask3"].upstream_task_ids == {"mytask2", "setuptask2"} + assert dag.task_group.children["mytask3"].downstream_task_ids == {"mytask4"} + assert dag.task_group.children["mytask4"].upstream_task_ids == {"mytask3"} Review Comment: ```suggestion assert dag.task_group.children["mytask4"].upstream_task_ids == {"mytask3"} assert not dag.task_group.children["mytask4"].downstream_task_ids ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
