code: def test_sub_dag_task_group():
      """
      Tests dag.partial_subset() updates task_group correctly.
      """
      execution_date = pendulum.parse("20200101")
      with DAG(
          "test_test_task_group_sub_dag", start_date=execution_date
      ) as dag:
          task1 = EmptyOperator(task_id="task1")
          with TaskGroup("group234") as group234:
              _ = EmptyOperator(task_id="task2")
  
              with TaskGroup("group34") as group34:
                  _ = EmptyOperator(task_id="task3")
                  _ = EmptyOperator(task_id="task4")
  
          with TaskGroup("group6") as group6:
              _ = EmptyOperator(task_id="task6")
  
          task7 = EmptyOperator(task_id="task7")
          task5 = EmptyOperator(task_id="task5")
  
          task1 >> group234
          group34 >> task5
          group234 >> group6
          group234 >> task7
  
      subdag = dag.partial_subset(
          task_ids_or_regex="task5",
          include_upstream=True,
          include_downstream=False,
      )
  
      assert extract_node_id(task_group_to_dict(subdag.task_group)) == {
          "id": None,
          "children": [
              {
                  "id": "group234",
                  "children": [
                      {
                          "id": "group234.group34",
                          "children": [
                              {"id": "group234.group34.task3"},
                              {"id": "group234.group34.task4"},
                              {
                                  "id": "group234.group34.downstream_join_id"
                              },
                          ],
                      },
                      {"id": "group234.upstream_join_id"},
                  ],
              },
              {"id": "task1"},
              {"id": "task5"},
          ],
      }
  
      edges = dag_edges(subdag)
      assert sorted(
          (e["source_id"], e["target_id"]) for e in edges
      ) == [
          ("group234.group34.downstream_join_id", "task5"),
          (
              "group234.group34.task3",
              "group234.group34.downstream_join_id",
          ),
          (
              "group234.group34.task4",
              "group234.group34.downstream_join_id",
          ),
          ("group234.upstream_join_id", "group234.group34.task3"),
          ("group234.upstream_join_id", "group234.group34.task4"),
          ("task1", "group234.upstream_join_id"),
      ]
  
      subdag_task_groups = subdag.task_group.get_task_group_dict()
      assert subdag_task_groups.keys() == {
          None,
          "group234",
          "group234.group34",
      }
  
      included_group_ids = {"group234", "group234.group34"}
      included_task_ids = {
          "group234.group34.task3",
          "group234.group34.task4",
          "task1",
          "task5",
      }
  
      for task_group in subdag_task_groups.values():
          assert task_group.upstream_group_ids.issubset(
              included_group_ids
          )
          assert task_group.downstream_group_ids.issubset(
              included_group_ids
          )
          assert task_group.upstream_task_ids.issubset(
              included_task_ids
          )
          assert task_group.downstream_task_ids.issubset(
              included_task_ids
          )
  
      for task in subdag.task_group:
          assert task.upstream_task_ids.issubset(included_task_ids)
          assert task.downstream_task_ids.issubset(included_task_ids)
  
imports: from __future__ import annotations
  from datetime import timedelta
  import pendulum
  import pytest
  from airflow.decorators import dag, setup, task, task_group, teardown
  from airflow.exceptions import TaskAlreadyInTaskGroup
  from airflow.models.baseoperator import BaseOperator
  from airflow.models.dag import DAG
  from airflow.models.xcom_arg import XComArg
  from airflow.operators.bash import BashOperator
  from airflow.operators.empty import EmptyOperator
  from airflow.operators.python import PythonOperator
  from airflow.utils.dag_edges import dag_edges
  from airflow.utils.task_group import TaskGroup, task_group_to_dict
  from tests.models import DEFAULT_DATE
label_desc: Use image node_id 7 for calling a global function with example usage: test_sub_dag_task_group() without return types