test cases for setting dependencies correctly between task group and other if task group is empty #47947

Author: purnachandergit

airflow-core/tests/unit/utils/test_task_group.py

@@ -1737,3 +1737,69 @@ def test_task_group_display_name_used_as_label():
     }
 
     assert extract_node_id(task_group_to_dict(dag.task_group), include_label=True) == expected_node_id
+
+
+def test_task_group_dependencies_between_tasks_if_task_group_is_empty_1():
+    """
+    Test that if a task group is empty, the dependencies between tasks are still maintained.
+    """
+    with DAG(dag_id="test_dag", schedule=None, start_date=pendulum.parse("20200101")):
+        task1 = EmptyOperator(task_id="task1")
+        with TaskGroup("group1") as tg1:
+            pass
+        with TaskGroup("group2") as tg2:
+            task2 = EmptyOperator(task_id="task2")
+            task3 = EmptyOperator(task_id="task3")
+            task2 >> task3
+
+        task1 >> tg1 >> tg2
+
+    assert task1.downstream_task_ids == {"group2.task2"}
+
+
+def test_task_group_dependencies_between_tasks_if_task_group_is_empty_2():
+    """
+    Test that if a task group is empty, the dependencies between tasks are still maintained.
+    """
+    with DAG(dag_id="test_dag", schedule=None, start_date=pendulum.parse("20200101")):
+        task1 = EmptyOperator(task_id="task1")
+        with TaskGroup("group1") as tg1:
+            pass
+        with TaskGroup("group2") as tg2:
+            pass
+        with TaskGroup("group3") as tg3:
+            pass
+        with TaskGroup("group4") as tg4:
+            pass
+        with TaskGroup("group5") as tg5:
+            task2 = EmptyOperator(task_id="task2")
+            task3 = EmptyOperator(task_id="task3")
+            task2 >> task3
+        task1 >> tg1 >> tg2 >> tg3 >> tg4 >> tg5
+
+    assert task1.downstream_task_ids == {"group5.task2"}
+
+
+def test_task_group_dependencies_between_tasks_if_task_group_is_empty_3():
+    """
+    Test that if a task group is empty, the dependencies between tasks are still maintained.
+    """
+    with DAG(dag_id="test_dag", schedule=None, start_date=pendulum.parse("20200101")):
+        task1 = EmptyOperator(task_id="task1")
+        with TaskGroup("group1") as tg1:
+            pass
+        with TaskGroup("group2") as tg2:
+            pass
+        task2 = EmptyOperator(task_id="task2")
+        with TaskGroup("group3") as tg3:
+            pass
+        with TaskGroup("group4") as tg4:
+            pass
+        with TaskGroup("group5") as tg5:
+            task3 = EmptyOperator(task_id="task3")
+            task4 = EmptyOperator(task_id="task4")
+            task3 >> task4
+        task1 >> tg1 >> tg2 >> task2 >> tg3 >> tg4 >> tg5
+
+    assert task1.downstream_task_ids == {"task2"}
+    assert task2.downstream_task_ids == {"group5.task3"}

task-sdk/src/airflow/sdk/definitions/taskgroup.py

@@ -344,10 +344,11 @@ def _set_relatives(
             task_or_task_list = [task_or_task_list]
 
         # Helper function to find leaves from a task list or task group
-        def find_leaves(group_or_task) -> list:
+        def find_leaves(group_or_task) -> list[Any]:
             while group_or_task:
-                if list(group_or_task.get_leaves()):
-                    return list(group_or_task.get_leaves())
+                group_or_task_leaves = list(group_or_task.get_leaves())
+                if group_or_task_leaves:
+                    return group_or_task_leaves
                 if group_or_task.upstream_task_ids:
                     upstream_task_ids_list = list(group_or_task.upstream_task_ids)
                     return [self.dag.get_task(task_id) for task_id in upstream_task_ids_list]