yuqian90 commented on a change in pull request #10153: URL: https://github.com/apache/airflow/pull/10153#discussion_r485543331
########## File path: tests/utils/test_task_group.py ########## @@ -0,0 +1,523 @@ +# +# 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. + +import pendulum +import pytest + +from airflow.models import DAG +from airflow.operators.dummy_operator import DummyOperator +from airflow.utils.task_group import TaskGroup +from airflow.www.views import dag_edges, task_group_to_dict + +EXPECTED_JSON = { + 'id': None, + 'value': { + 'label': None, + 'labelStyle': 'fill:#000;', + 'style': 'fill:CornflowerBlue', + 'rx': 5, + 'ry': 5, + 'clusterLabelPos': 'top', + }, + 'tooltip': '', + 'children': [ + { + 'id': 'group234', + 'value': { + 'label': 'group234', + 'labelStyle': 'fill:#000;', + 'style': 'fill:CornflowerBlue', + 'rx': 5, + 'ry': 5, + 'clusterLabelPos': 'top', + }, + 'tooltip': '', + 'children': [ + { + 'id': 'group234.group34', + 'value': { + 'label': 'group34', + 'labelStyle': 'fill:#000;', + 'style': 'fill:CornflowerBlue', + 'rx': 5, + 'ry': 5, + 'clusterLabelPos': 'top', + }, + 'tooltip': '', + 'children': [ + { + 'id': 'group234.group34.task3', + 'value': { + 'label': 'task3', + 'labelStyle': 'fill:#000;', + 'style': 'fill:#e8f7e4;', + 'rx': 5, + 'ry': 5, + }, + }, + { + 'id': 'group234.group34.task4', + 'value': { + 'label': 'task4', + 'labelStyle': 'fill:#000;', + 'style': 'fill:#e8f7e4;', + 'rx': 5, + 'ry': 5, + }, + }, + { + 'id': 'group234.group34.downstream_join_id', + 'value': { + 'label': '', + 'labelStyle': 'fill:#000;', + 'style': 'fill:CornflowerBlue;', + 'shape': 'circle', + }, + }, + ], + }, + { + 'id': 'group234.task2', + 'value': { + 'label': 'task2', + 'labelStyle': 'fill:#000;', + 'style': 'fill:#e8f7e4;', + 'rx': 5, + 'ry': 5, + }, + }, + { + 'id': 'group234.upstream_join_id', + 'value': { + 'label': '', + 'labelStyle': 'fill:#000;', + 'style': 'fill:CornflowerBlue;', + 'shape': 'circle', + }, + }, + ], + }, + { + 'id': 'task1', + 'value': { + 'label': 'task1', + 'labelStyle': 'fill:#000;', + 'style': 'fill:#e8f7e4;', + 'rx': 5, + 'ry': 5, + }, + }, + { + 'id': 'task5', + 'value': { + 'label': 'task5', + 'labelStyle': 'fill:#000;', + 'style': 'fill:#e8f7e4;', + 'rx': 5, + 'ry': 5, + }, + }, + ], +} + + +def test_build_task_group_context_manager(): + execution_date = pendulum.parse("20200101") + with DAG("test_build_task_group_context_manager", start_date=execution_date) as dag: + task1 = DummyOperator(task_id="task1") + with TaskGroup("group234") as group234: + _ = DummyOperator(task_id="task2") + + with TaskGroup("group34") as group34: + _ = DummyOperator(task_id="task3") + _ = DummyOperator(task_id="task4") + + task5 = DummyOperator(task_id="task5") + task1 >> group234 + group34 >> task5 + + assert task1.get_direct_relative_ids(upstream=False) == { + 'group234.group34.task4', + 'group234.group34.task3', + 'group234.task2', + } + assert task5.get_direct_relative_ids(upstream=True) == { + 'group234.group34.task4', + 'group234.group34.task3', + } + + assert dag.task_group.group_id is None + assert dag.task_group.is_root + assert set(dag.task_group.children.keys()) == {"task1", "group234", "task5"} + assert group34.group_id == "group234.group34" + + assert task_group_to_dict(dag.task_group) == EXPECTED_JSON + + +def test_build_task_group(): + """ + This is an alternative syntax to use TaskGroup. It should result in the same TaskGroup + as using context manager. + """ + execution_date = pendulum.parse("20200101") + dag = DAG("test_build_task_group", start_date=execution_date) + task1 = DummyOperator(task_id="task1", dag=dag) + group234 = TaskGroup("group234", dag=dag) + _ = DummyOperator(task_id="task2", dag=dag, task_group=group234) + group34 = TaskGroup("group34", dag=dag, parent_group=group234) + _ = DummyOperator(task_id="task3", dag=dag, task_group=group34) + _ = DummyOperator(task_id="task4", dag=dag, task_group=group34) + task5 = DummyOperator(task_id="task5", dag=dag) + + task1 >> group234 + group34 >> task5 + + assert task_group_to_dict(dag.task_group) == EXPECTED_JSON + + +def extract_node_id(node, include_label=False): + ret = {"id": node["id"]} + if include_label: + ret["label"] = node["value"]["label"] + if "children" in node: + children = [] + for child in node["children"]: + children.append(extract_node_id(child, include_label=include_label)) + + ret["children"] = children + + return ret + + +def test_build_task_group_with_prefix(): + """ + Tests that prefix_group_id turns on/off prefixing of task_id with group_id. + """ + execution_date = pendulum.parse("20200101") + with DAG("test_build_task_group_with_prefix", start_date=execution_date) as dag: + task1 = DummyOperator(task_id="task1") + with TaskGroup("group234", prefix_group_id=False) as group234: + task2 = DummyOperator(task_id="task2") + + with TaskGroup("group34") as group34: + task3 = DummyOperator(task_id="task3") + + with TaskGroup("group4", prefix_group_id=False) as group4: + task4 = DummyOperator(task_id="task4") + + task5 = DummyOperator(task_id="task5") + task1 >> group234 >> task5 + + assert task2.task_id == "task2" + assert group34.group_id == "group34" + assert task3.task_id == "group34.task3" + assert group4.group_id == "group34.group4" + assert task4.task_id == "task4" + assert task5.task_id == "task5" + assert group234.get_child_by_label("task2") == task2 + assert group234.get_child_by_label("group34") == group34 + assert group4.get_child_by_label("task4") == task4 + + assert extract_node_id(task_group_to_dict(dag.task_group), include_label=True) == { + 'id': None, + 'label': None, + 'children': [ + { + 'id': 'group234', + 'label': 'group234', + 'children': [ + { + 'id': 'group34', + 'label': 'group34', + 'children': [ + { + 'id': 'group34.group4', + 'label': 'group4', + 'children': [{'id': 'task4', 'label': 'task4'}], + }, + {'id': 'group34.task3', 'label': 'task3'}, + {'id': 'group34.downstream_join_id', 'label': ''}, + ], + }, + {'id': 'task2', 'label': 'task2'}, + {'id': 'group234.upstream_join_id', 'label': ''}, + ], + }, + {'id': 'task1', 'label': 'task1'}, + {'id': 'task5', 'label': 'task5'}, + ], + } + + +def test_build_task_group_with_task_decorator(): + """ + Test that TaskGroup can be used with the @task decorator. + """ + from airflow.operators.python import task + + @task + def task_1(): + print("task_1") + + @task + def task_2(): + return "task_2" + + @task + def task_3(): + return "task_3" + + @task + def task_4(task_2_output, task_3_output): + print(task_2_output, task_3_output) + + @task + def task_5(): + print("task_5") + + execution_date = pendulum.parse("20200101") + with DAG("test_build_task_group_with_task_decorator", start_date=execution_date) as dag: + tsk_1 = task_1() + + with TaskGroup("group234") as group234: + tsk_2 = task_2() + tsk_3 = task_3() + _ = task_4(tsk_2, tsk_3) + + tsk_5 = task_5() + + tsk_1 >> group234 + tsk_5 << group234 Review comment: Sure that's done. ---------------------------------------------------------------- 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. For queries about this service, please contact Infrastructure at: [email protected]
