dstandish commented on code in PR #30342:
URL: https://github.com/apache/airflow/pull/30342#discussion_r1150810313


##########
tests/decorators/test_setup_teardown.py:
##########
@@ -153,3 +156,57 @@ def 
test_classic_teardown_task_works_with_on_failure_fail_dagrun(self, on_failur
         assert teardown_task._is_teardown
         assert teardown_task._on_failure_fail_dagrun is on_failure_fail_dagrun
         assert len(dag.task_group.children) == 1
+
+    def test_setup_task_can_be_overriden(self, dag_maker):
+        @setup
+        def mytask():
+            print("I am a setup task")
+
+        with dag_maker() as dag:
+            mytask.override(task_id="mytask2")()
+        assert len(dag.task_group.children) == 1
+        setup_task = dag.task_group.children["mytask2"]
+        assert setup_task._is_setup
+
+    def test_setup_teardown_mixed_up_in_a_dag(self, dag_maker):
+        @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:
+            setuptask()
+            teardowntask()
+            setuptask2()
+            teardowntask2()
+            mytask()
+            mytask2()
+
+        assert len(dag.task_group.children) == 6

Review Comment:
   ```suggestion
           assert len(dag.task_group.children) == 6
           assert [x for x in dag.tasks if not x.downstream_list]  # no deps 
have been set
   
   ```



-- 
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]

Reply via email to