| # |
| # 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 collections.abc import Callable |
| from datetime import datetime, timedelta |
| from typing import TYPE_CHECKING |
| from unittest import mock |
| |
| import pendulum |
| import pytest |
| |
| from airflow.sdk import TaskInstanceState, TriggerRule |
| from airflow.sdk.bases.operator import BaseOperator |
| from airflow.sdk.bases.xcom import BaseXCom |
| from airflow.sdk.definitions.dag import DAG |
| from airflow.sdk.definitions.mappedoperator import MappedOperator |
| from airflow.sdk.definitions.xcom_arg import XComArg |
| from airflow.sdk.execution_time.comms import ( |
| GetTICount, |
| GetXCom, |
| GetXComSequenceSlice, |
| SetXCom, |
| TICount, |
| XComResult, |
| XComSequenceSliceResult, |
| ) |
| |
| from tests_common.test_utils.mapping import expand_mapped_task # noqa: F401 |
| from tests_common.test_utils.mock_operators import ( |
| MockOperator, |
| ) |
| |
| DEFAULT_DATE = datetime(2016, 1, 1) |
| |
| |
| def test_task_mapping_with_dag(): |
| with DAG("test-dag") as dag: |
| task1 = BaseOperator(task_id="op1") |
| literal = ["a", "b", "c"] |
| mapped = MockOperator.partial(task_id="task_2").expand(arg2=literal) |
| finish = MockOperator(task_id="finish") |
| |
| task1 >> mapped >> finish |
| |
| assert task1.downstream_list == [mapped] |
| assert mapped in dag.tasks |
| assert mapped.task_group == dag.task_group |
| # At parse time there should only be three tasks! |
| assert len(dag.tasks) == 3 |
| |
| assert finish.upstream_list == [mapped] |
| assert mapped.downstream_list == [finish] |
| |
| |
| # TODO: |
| # test_task_mapping_with_dag_and_list_of_pandas_dataframe |
| |
| |
| def test_task_mapping_without_dag_context(): |
| with DAG("test-dag") as dag: |
| task1 = BaseOperator(task_id="op1") |
| literal = ["a", "b", "c"] |
| mapped = MockOperator.partial(task_id="task_2").expand(arg2=literal) |
| |
| task1 >> mapped |
| |
| assert isinstance(mapped, MappedOperator) |
| assert mapped in dag.tasks |
| assert task1.downstream_list == [mapped] |
| assert mapped in dag.tasks |
| # At parse time there should only be two tasks! |
| assert len(dag.tasks) == 2 |
| |
| |
| def test_task_mapping_default_args(): |
| default_args = {"start_date": DEFAULT_DATE.now(), "owner": "test"} |
| with DAG("test-dag", schedule=None, start_date=DEFAULT_DATE, default_args=default_args): |
| task1 = BaseOperator(task_id="op1") |
| literal = ["a", "b", "c"] |
| mapped = MockOperator.partial(task_id="task_2").expand(arg2=literal) |
| |
| task1 >> mapped |
| |
| assert mapped.partial_kwargs["owner"] == "test" |
| assert mapped.start_date == pendulum.instance(default_args["start_date"]) |
| |
| |
| def test_task_mapping_override_default_args(): |
| default_args = {"retries": 2, "start_date": DEFAULT_DATE.now()} |
| with DAG("test-dag", schedule=None, start_date=DEFAULT_DATE, default_args=default_args): |
| literal = ["a", "b", "c"] |
| mapped = MockOperator.partial(task_id="task", retries=1).expand(arg2=literal) |
| |
| # retries should be 1 because it is provided as a partial arg |
| assert mapped.partial_kwargs["retries"] == 1 |
| # start_date should be equal to default_args["start_date"] because it is not provided as partial arg |
| assert mapped.start_date == pendulum.instance(default_args["start_date"]) |
| # owner should be equal to Airflow default owner (airflow) because it is not provided at all |
| assert mapped.owner == "airflow" |
| |
| |
| def test_map_unknown_arg_raises(): |
| with pytest.raises(TypeError, match=r"argument 'file'"): |
| BaseOperator.partial(task_id="a").expand(file=[1, 2, {"a": "b"}]) |
| |
| |
| def test_map_xcom_arg(): |
| """Test that dependencies are correct when mapping with an XComArg""" |
| with DAG("test-dag"): |
| task1 = BaseOperator(task_id="op1") |
| mapped = MockOperator.partial(task_id="task_2").expand(arg2=task1.output) |
| finish = MockOperator(task_id="finish") |
| |
| mapped >> finish |
| |
| assert task1.downstream_list == [mapped] |
| |
| |
| def test_partial_on_instance() -> None: |
| """`.partial` on an instance should fail -- it's only designed to be called on classes""" |
| with pytest.raises(TypeError): |
| MockOperator(task_id="a").partial() |
| |
| |
| def test_partial_on_class() -> None: |
| # Test that we accept args for superclasses too |
| op = MockOperator.partial( |
| task_id="a", |
| arg1="a", |
| trigger_rule=TriggerRule.ONE_FAILED, |
| on_execute_callback=id, |
| ) |
| assert op.kwargs["arg1"] == "a" |
| assert op.kwargs["trigger_rule"] == TriggerRule.ONE_FAILED |
| assert op.kwargs["on_execute_callback"] == [id] |
| |
| |
| def test_partial_on_class_invalid_ctor_args() -> None: |
| """Test that when we pass invalid args to partial(). |
| |
| I.e. if an arg is not known on the class or any of its parent classes we error at parse time |
| """ |
| with pytest.raises(TypeError, match=r"arguments 'foo', 'bar'"): |
| MockOperator.partial(task_id="a", foo="bar", bar=2) |
| |
| |
| def test_partial_on_invalid_pool_slots_raises() -> None: |
| """Test that when we pass an invalid value to pool_slots in partial(), |
| |
| i.e. if the value is not an integer, an error is raised at import time.""" |
| |
| with pytest.raises(TypeError, match="'<' not supported between instances of 'str' and 'int'"): |
| MockOperator.partial(task_id="pool_slots_test", pool="test", pool_slots="a").expand(arg1=[1, 2, 3]) |
| |
| |
| def test_mapped_task_applies_default_args_classic(): |
| with DAG( |
| dag_id="test", |
| default_args={"execution_timeout": timedelta(minutes=30), "on_failure_callback": str}, |
| ) as dag: |
| MockOperator(task_id="simple", arg1=None, arg2=0) |
| MockOperator.partial(task_id="mapped").expand(arg1=[1], arg2=[2, 3]) |
| |
| assert dag.get_task("simple").execution_timeout == timedelta(minutes=30) |
| assert dag.get_task("mapped").execution_timeout == timedelta(minutes=30) |
| assert dag.get_task("simple").on_failure_callback == [str] |
| assert dag.get_task("mapped").on_failure_callback == [str] |
| |
| |
| def test_mapped_task_applies_default_args_taskflow(): |
| with DAG("test", default_args={"execution_timeout": timedelta(minutes=30)}) as dag: |
| |
| @dag.task |
| def simple(arg): |
| pass |
| |
| @dag.task |
| def mapped(arg): |
| pass |
| |
| simple(arg=0) |
| mapped.expand(arg=[1, 2]) |
| |
| assert dag.get_task("simple").execution_timeout == timedelta(minutes=30) |
| assert dag.get_task("mapped").execution_timeout == timedelta(minutes=30) |
| assert dag.get_task("simple").on_success_callback == [] |
| assert dag.get_task("mapped").on_success_callback == [] |
| |
| |
| @pytest.mark.parametrize( |
| ("callable", "expected"), |
| [ |
| pytest.param( |
| lambda partial, output1: partial.expand( |
| map_template=output1, map_static=output1, file_template=["/path/to/file.ext"] |
| ), |
| # Note to the next person to come across this. In #32272 we changed expand_kwargs so that it |
| # resolves the mapped template when it's in `expand_kwargs()`, but we _didn't_ do the same for |
| # things in `expand()`. This feels like a bug to me (ashb) but I am not changing that now, I have |
| # just moved and parametrized this test. |
| "{{ ds }}", |
| id="expand", |
| ), |
| pytest.param( |
| lambda partial, output1: partial.expand_kwargs( |
| [{"map_template": "{{ ds }}", "map_static": "{{ ds }}", "file_template": "/path/to/file.ext"}] |
| ), |
| "2024-12-01", |
| id="expand_kwargs", |
| ), |
| ], |
| ) |
| def test_mapped_render_template_fields_validating_operator( |
| tmp_path, create_runtime_ti, mock_supervisor_comms, callable, expected: bool |
| ): |
| file_template_dir = tmp_path / "path" / "to" |
| file_template_dir.mkdir(parents=True, exist_ok=True) |
| file_template = file_template_dir / "file.ext" |
| file_template.write_text("loaded data") |
| |
| class MyOperator(BaseOperator): |
| template_fields = ("partial_template", "map_template", "file_template") |
| template_ext = (".ext",) |
| |
| def __init__( |
| self, partial_template, partial_static, map_template, map_static, file_template, **kwargs |
| ): |
| for value in [partial_template, partial_static, map_template, map_static, file_template]: |
| assert isinstance(value, str), "value should have been resolved before unmapping" |
| super().__init__(**kwargs) |
| self.partial_template = partial_template |
| self.partial_static = partial_static |
| self.map_template = map_template |
| self.map_static = map_static |
| self.file_template = file_template |
| |
| def execute(self, context): |
| pass |
| |
| with DAG("test_dag", template_searchpath=tmp_path.__fspath__()): |
| task1 = BaseOperator(task_id="op1") |
| mapped = MyOperator.partial( |
| task_id="a", partial_template="{{ ti.task_id }}", partial_static="{{ ti.task_id }}" |
| ) |
| mapped = callable(mapped, task1.output) |
| |
| def mock_comms(msg): |
| if isinstance(msg, GetXCom): |
| return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=["{{ ds }}"]) |
| if isinstance(msg, GetXComSequenceSlice): |
| return XComSequenceSliceResult(root=["{{ ds }}"]) |
| if isinstance(msg, GetTICount): |
| return TICount(count=1) |
| return mock.DEFAULT |
| |
| mock_supervisor_comms.send.side_effect = mock_comms |
| |
| mapped_ti = create_runtime_ti(task=mapped, map_index=0) |
| |
| assert isinstance(mapped_ti.task, MappedOperator) |
| mapped_ti.task.render_template_fields(context=mapped_ti.get_template_context()) |
| assert isinstance(mapped_ti.task, MyOperator) |
| |
| assert mapped_ti.task.partial_template == "a", "Should be templated!" |
| assert mapped_ti.task.partial_static == "{{ ti.task_id }}", "Should not be templated!" |
| assert mapped_ti.task.map_template == expected |
| assert mapped_ti.task.map_static == "{{ ds }}", "Should not be templated!" |
| assert mapped_ti.task.file_template == "loaded data", "Should be templated!" |
| |
| |
| def test_mapped_render_nested_template_fields(create_runtime_ti, mock_supervisor_comms): |
| with DAG("test_dag"): |
| mapped = MockOperatorWithNestedFields.partial( |
| task_id="t", arg2=NestedFields(field_1="{{ ti.task_id }}", field_2="value_2") |
| ).expand(arg1=["{{ ti.task_id }}", ["s", "{{ ti.task_id }}"]]) |
| |
| ti = create_runtime_ti(task=mapped, map_index=0) |
| ti.task.render_template_fields(context=ti.get_template_context()) |
| assert ti.task.arg1 == "t" |
| assert ti.task.arg2.field_1 == "t" |
| assert ti.task.arg2.field_2 == "value_2" |
| |
| ti = create_runtime_ti(task=mapped, map_index=1) |
| ti.task.render_template_fields(context=ti.get_template_context()) |
| assert ti.task.arg1 == ["s", "t"] |
| assert ti.task.arg2.field_1 == "t" |
| assert ti.task.arg2.field_2 == "value_2" |
| |
| |
| @pytest.mark.parametrize( |
| ("map_index", "expected"), |
| [ |
| pytest.param(0, "2024-12-01", id="0"), |
| pytest.param(1, 2, id="1"), |
| ], |
| ) |
| def test_expand_kwargs_render_template_fields_validating_operator( |
| map_index, expected, create_runtime_ti, mock_supervisor_comms |
| ): |
| with DAG("test_dag"): |
| task1 = BaseOperator(task_id="op1") |
| mapped = MockOperator.partial(task_id="a", arg2="{{ ti.task_id }}").expand_kwargs(task1.output) |
| |
| xcom_values = [{"arg1": "{{ ds }}"}, {"arg1": 2}] |
| |
| def mock_comms(msg): |
| if isinstance(msg, GetXCom): |
| return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=xcom_values) |
| if isinstance(msg, GetXComSequenceSlice): |
| return XComSequenceSliceResult(root=xcom_values) |
| if isinstance(msg, GetTICount): |
| return TICount(count=2) |
| return mock.DEFAULT |
| |
| mock_supervisor_comms.send.side_effect = mock_comms |
| |
| ti = create_runtime_ti(task=mapped, map_index=map_index) |
| assert isinstance(ti.task, MappedOperator) |
| ti.task.render_template_fields(context=ti.get_template_context()) |
| assert isinstance(ti.task, MockOperator) |
| assert ti.task.arg1 == expected |
| assert ti.task.arg2 == "a" |
| |
| |
| def test_xcomarg_property_of_mapped_operator(): |
| with DAG("test_xcomarg_property_of_mapped_operator"): |
| op_a = MockOperator.partial(task_id="a").expand(arg1=["x", "y", "z"]) |
| |
| assert op_a.output == XComArg(op_a) |
| |
| |
| def test_set_xcomarg_dependencies_with_mapped_operator(): |
| with DAG("test_set_xcomargs_dependencies_with_mapped_operator"): |
| op1 = MockOperator.partial(task_id="op1").expand(arg1=[1, 2, 3]) |
| op2 = MockOperator.partial(task_id="op2").expand(arg2=["a", "b", "c"]) |
| op3 = MockOperator(task_id="op3", arg1=op1.output) |
| op4 = MockOperator(task_id="op4", arg1=[op1.output, op2.output]) |
| op5 = MockOperator(task_id="op5", arg1={"op1": op1.output, "op2": op2.output}) |
| |
| assert op1 in op3.upstream_list |
| assert op1 in op4.upstream_list |
| assert op2 in op4.upstream_list |
| assert op1 in op5.upstream_list |
| assert op2 in op5.upstream_list |
| |
| |
| def test_task_mapping_with_task_group_context(): |
| from airflow.sdk.definitions.taskgroup import TaskGroup |
| |
| with DAG("test-dag") as dag: |
| task1 = BaseOperator(task_id="op1") |
| finish = MockOperator(task_id="finish") |
| |
| with TaskGroup("test-group") as group: |
| literal = ["a", "b", "c"] |
| mapped = MockOperator.partial(task_id="task_2").expand(arg2=literal) |
| |
| task1 >> group >> finish |
| |
| assert task1.downstream_list == [mapped] |
| assert mapped.upstream_list == [task1] |
| |
| assert mapped in dag.tasks |
| assert mapped.task_group == group |
| |
| assert finish.upstream_list == [mapped] |
| assert mapped.downstream_list == [finish] |
| |
| |
| def test_task_mapping_with_explicit_task_group(): |
| from airflow.sdk.definitions.taskgroup import TaskGroup |
| |
| with DAG("test-dag") as dag: |
| task1 = BaseOperator(task_id="op1") |
| finish = MockOperator(task_id="finish") |
| |
| group = TaskGroup("test-group") |
| literal = ["a", "b", "c"] |
| mapped = MockOperator.partial(task_id="task_2", task_group=group).expand(arg2=literal) |
| |
| task1 >> group >> finish |
| |
| assert task1.downstream_list == [mapped] |
| assert mapped.upstream_list == [task1] |
| |
| assert mapped in dag.tasks |
| assert mapped.task_group == group |
| |
| assert finish.upstream_list == [mapped] |
| assert mapped.downstream_list == [finish] |
| |
| |
| def test_nested_mapped_task_groups(): |
| from airflow.decorators import task, task_group |
| |
| with DAG("test"): |
| |
| @task |
| def t(): |
| return [[1, 2], [3, 4]] |
| |
| @task |
| def m(x): |
| return x |
| |
| @task_group |
| def g1(x): |
| @task_group |
| def g2(y): |
| return m(y) |
| |
| return g2.expand(y=x) |
| |
| # Add a test once nested mapped task groups become supported |
| with pytest.raises(ValueError, match="Nested Mapped TaskGroups are not yet supported"): |
| g1.expand(x=t()) |
| |
| |
| RunTI = Callable[[DAG, str, int], TaskInstanceState] |
| |
| |
| def test_map_cross_product(run_ti: RunTI, mock_supervisor_comms): |
| outputs = [] |
| |
| with DAG(dag_id="cross_product") as dag: |
| |
| @dag.task |
| def emit_numbers(): |
| return [1, 2] |
| |
| @dag.task |
| def emit_letters(): |
| return {"a": "x", "b": "y", "c": "z"} |
| |
| @dag.task |
| def show(number, letter): |
| outputs.append((number, letter)) |
| |
| show.expand(number=emit_numbers(), letter=emit_letters()) |
| |
| numbers = [1, 2] |
| letters = {"a": "x", "b": "y", "c": "z"} |
| |
| def mock_comms(msg): |
| if isinstance(msg, GetXCom): |
| if msg.task_id == "emit_numbers": |
| return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=numbers) |
| if msg.task_id == "emit_letters": |
| return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=letters) |
| elif isinstance(msg, GetXComSequenceSlice): |
| if msg.task_id == "emit_numbers": |
| return XComSequenceSliceResult(root=numbers) |
| if msg.task_id == "emit_letters": |
| # Convert dict items to list for XComSequenceSliceResult |
| return XComSequenceSliceResult(root=list(letters.items())) |
| elif isinstance(msg, GetTICount): |
| # show is mapped by 6 (2 numbers * 3 letters) |
| if msg.task_ids and msg.task_ids[0] == "show": |
| return TICount(count=6) |
| return TICount(count=1) |
| return mock.DEFAULT |
| |
| mock_supervisor_comms.send.side_effect = mock_comms |
| |
| states = [run_ti(dag, "show", map_index) for map_index in range(6)] |
| assert states == [TaskInstanceState.SUCCESS] * 6 |
| assert outputs == [ |
| (1, ("a", "x")), |
| (1, ("b", "y")), |
| (1, ("c", "z")), |
| (2, ("a", "x")), |
| (2, ("b", "y")), |
| (2, ("c", "z")), |
| ] |
| |
| |
| def test_map_product_same(run_ti: RunTI, mock_supervisor_comms): |
| """Test a mapped task can refer to the same source multiple times.""" |
| outputs = [] |
| |
| with DAG(dag_id="product_same") as dag: |
| |
| @dag.task |
| def emit_numbers(): |
| return [1, 2] |
| |
| @dag.task |
| def show(a, b): |
| outputs.append((a, b)) |
| |
| emit_task = emit_numbers() |
| show.expand(a=emit_task, b=emit_task) |
| |
| numbers = [1, 2] |
| |
| def mock_comms(msg): |
| if isinstance(msg, GetXCom): |
| if msg.task_id == "emit_numbers": |
| return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=numbers) |
| elif isinstance(msg, GetXComSequenceSlice): |
| if msg.task_id == "emit_numbers": |
| return XComSequenceSliceResult(root=numbers) |
| elif isinstance(msg, GetTICount): |
| # show is mapped by 4 (2 * 2 cross product) |
| if msg.task_ids and msg.task_ids[0] == "show": |
| return TICount(count=4) |
| return TICount(count=1) |
| return mock.DEFAULT |
| |
| mock_supervisor_comms.send.side_effect = mock_comms |
| |
| states = [run_ti(dag, "show", map_index) for map_index in range(4)] |
| assert states == [TaskInstanceState.SUCCESS] * 4 |
| assert outputs == [(1, 1), (1, 2), (2, 1), (2, 2)] |
| |
| |
| class NestedFields: |
| """Nested fields for testing purposes.""" |
| |
| def __init__(self, field_1, field_2): |
| self.field_1 = field_1 |
| self.field_2 = field_2 |
| |
| |
| class MockOperatorWithNestedFields(BaseOperator): |
| """Operator with nested fields for testing purposes.""" |
| |
| template_fields = ("arg1", "arg2") |
| |
| def __init__(self, arg1: str = "", arg2: NestedFields | None = None, **kwargs): |
| super().__init__(**kwargs) |
| self.arg1 = arg1 |
| self.arg2 = arg2 |
| |
| def _render_nested_template_fields(self, content, context, jinja_env, seen_oids) -> None: |
| if id(content) not in seen_oids: |
| template_fields: tuple | None = None |
| |
| if isinstance(content, NestedFields): |
| template_fields = ("field_1", "field_2") |
| |
| if template_fields: |
| seen_oids.add(id(content)) |
| self._do_render_template_fields(content, template_fields, context, jinja_env, seen_oids) |
| return |
| |
| super()._render_nested_template_fields(content, context, jinja_env, seen_oids) |
| |
| |
| def test_find_mapped_dependants_in_another_group(): |
| from airflow.decorators import task as task_decorator |
| from airflow.sdk import TaskGroup |
| |
| @task_decorator |
| def gen(x): |
| return list(range(x)) |
| |
| @task_decorator |
| def add(x, y): |
| return x + y |
| |
| with DAG(dag_id="test"): |
| with TaskGroup(group_id="g1"): |
| gen_result = gen(3) |
| with TaskGroup(group_id="g2"): |
| add_result = add.partial(y=1).expand(x=gen_result) |
| |
| dependants = list(gen_result.operator.iter_mapped_dependants()) |
| assert dependants == [add_result.operator] |
| |
| |
| @pytest.mark.parametrize( |
| ("partial_params", "mapped_params", "expected"), |
| [ |
| pytest.param(None, [{"a": 1}], [{"a": 1}], id="simple"), |
| pytest.param({"b": 2}, [{"a": 1}], [{"a": 1, "b": 2}], id="merge"), |
| pytest.param({"b": 2}, [{"a": 1, "b": 3}], [{"a": 1, "b": 3}], id="override"), |
| pytest.param({"b": 2}, [{"a": 1, "b": 3}, {"b": 1}], [{"a": 1, "b": 3}, {"b": 1}], id="multiple"), |
| ], |
| ) |
| def test_mapped_expand_against_params(create_runtime_ti, partial_params, mapped_params, expected): |
| with DAG("test"): |
| task = BaseOperator.partial(task_id="t", params=partial_params).expand(params=mapped_params) |
| |
| for map_index, expansion in enumerate(expected): |
| mapped_ti = create_runtime_ti(task=task, map_index=map_index) |
| mapped_ti.task.render_template_fields(context=mapped_ti.get_template_context()) |
| assert mapped_ti.task.params == expansion |
| |
| |
| def test_operator_mapped_task_group_receives_value(create_runtime_ti, mock_supervisor_comms): |
| # Test the runtime expansion behaviour of mapped task groups + mapped operators |
| results = {} |
| |
| from airflow.decorators import task_group |
| |
| with DAG("test") as dag: |
| |
| @dag.task |
| def t(value, *, ti=None): |
| results[(ti.task_id, ti.map_index)] = value |
| return value |
| |
| @task_group |
| def tg(va): |
| # Each expanded group has one t1 and t2 each. |
| t1 = t.override(task_id="t1")(va) |
| t2 = t.override(task_id="t2")(t1) |
| |
| with pytest.raises(NotImplementedError) as ctx: |
| t.override(task_id="t4").expand(value=va) |
| assert str(ctx.value) == "operator expansion in an expanded task group is not yet supported" |
| |
| return t2 |
| |
| # The group is mapped by 3. |
| tg1 = tg.expand( |
| va=[ |
| ["a", "b"], |
| [4], |
| ["z"], |
| ] |
| ) |
| |
| # Aggregates results from task group. |
| t.override(task_id="t3")(tg1) |
| |
| # Map task group IDs to their expansion counts |
| task_group_expansion = {"tg": 3} |
| |
| def mock_comms_response(msg): |
| if isinstance(msg, GetXCom): |
| key = (msg.task_id, msg.map_index) |
| if key in expected_values: |
| value = expected_values[key] |
| return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=value) |
| if msg.map_index is None: |
| # Get all mapped XComValues for this ti |
| value = [v for k, v in expected_values.items() if k[0] == msg.task_id] |
| return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=value) |
| elif isinstance(msg, GetXComSequenceSlice): |
| # Handle sequence slicing for pulling all XCom values from mapped tasks |
| task_id = msg.task_id |
| values = [v for k, v in expected_values.items() if k[0] == task_id and k[1] is not None] |
| return XComSequenceSliceResult(root=values) |
| elif isinstance(msg, GetTICount): |
| # Handle TI count queries for upstream_map_indexes computation |
| if msg.task_ids: |
| task_id = msg.task_ids[0] |
| if task_id in expansion_per_task_id: |
| return TICount(count=len(list(expansion_per_task_id[task_id]))) |
| return TICount(count=1) |
| if msg.task_group_id: |
| return TICount(count=task_group_expansion.get(msg.task_group_id, 0)) |
| return TICount(count=0) |
| return mock.DEFAULT |
| |
| mock_supervisor_comms.send.side_effect = mock_comms_response |
| |
| expected_values = { |
| ("tg.t1", 0): ["a", "b"], |
| ("tg.t1", 1): [4], |
| ("tg.t1", 2): ["z"], |
| ("tg.t2", 0): ["a", "b"], |
| ("tg.t2", 1): [4], |
| ("tg.t2", 2): ["z"], |
| ("t3", None): [["a", "b"], [4], ["z"]], |
| } |
| |
| # We hard-code the number of expansions here as the server is in charge of that. |
| expansion_per_task_id = { |
| "tg.t1": range(3), |
| "tg.t2": range(3), |
| "t3": [None], |
| } |
| for task in dag.tasks: |
| for map_index in expansion_per_task_id[task.task_id]: |
| mapped_ti = create_runtime_ti( |
| task=task.prepare_for_execution(), |
| map_index=map_index, |
| ) |
| context = mapped_ti.get_template_context() |
| mapped_ti.task.render_template_fields(context) |
| mapped_ti.task.execute(context) |
| assert results == expected_values |
| |
| |
| def test_mapped_xcom_push_skipped_tasks(create_runtime_ti, mock_supervisor_comms): |
| from airflow.sdk import task_group |
| |
| if TYPE_CHECKING: |
| from airflow.providers.standard.operators.empty import EmptyOperator |
| from airflow.providers.standard.operators.python import ShortCircuitOperator |
| else: |
| ShortCircuitOperator = pytest.importorskip( |
| "airflow.providers.standard.operators.python" |
| ).ShortCircuitOperator |
| EmptyOperator = pytest.importorskip("airflow.providers.standard.operators.empty").EmptyOperator |
| |
| with DAG("test") as dag: |
| |
| @task_group |
| def group(x): |
| short_op_push_xcom = ShortCircuitOperator( |
| task_id="push_xcom_from_shortcircuit", |
| python_callable=lambda arg: arg % 2 == 0, |
| op_kwargs={"arg": x}, |
| ) |
| empty_task = EmptyOperator(task_id="empty_task") |
| short_op_push_xcom >> empty_task |
| |
| group.expand(x=[0, 1]) |
| |
| for task in dag.tasks: |
| for map_index in range(2): |
| ti = create_runtime_ti(task=task.prepare_for_execution(), map_index=map_index) |
| context = ti.get_template_context() |
| ti.task.render_template_fields(context) |
| ti.task.execute(context) |
| |
| assert ti |
| mock_supervisor_comms.send.assert_has_calls( |
| [ |
| mock.call( |
| msg=SetXCom( |
| key="skipmixin_key", |
| value={"skipped": ["group.empty_task"]}, |
| dag_id=ti.dag_id, |
| run_id=ti.run_id, |
| task_id="group.push_xcom_from_shortcircuit", |
| map_index=1, |
| ), |
| ), |
| ] |
| ) |
| |
| |
| @pytest.mark.parametrize( |
| ("setter_name", "old_value", "new_value"), |
| [ |
| ("owner", "old_owner", "new_owner"), |
| ("map_index_template", "old_mit", "new_mit"), |
| ("trigger_rule", TriggerRule.ALL_SUCCESS, TriggerRule.ALL_FAILED), |
| ("is_setup", True, False), |
| ("is_teardown", True, False), |
| ("depends_on_past", True, False), |
| ("ignore_first_depends_on_past", True, False), |
| ("wait_for_past_depends_before_skipping", True, False), |
| ("wait_for_downstream", True, False), |
| ("retries", 3, 5), |
| ("queue", "old_queue", "new_queue"), |
| ("pool", "old_pool", "new_pool"), |
| ("pool_slots", 1, 10), |
| ("execution_timeout", timedelta(minutes=5), timedelta(minutes=10)), |
| ("max_retry_delay", timedelta(minutes=5), timedelta(minutes=10)), |
| ("retry_delay", timedelta(minutes=5), timedelta(minutes=10)), |
| ("retry_exponential_backoff", 2.0, 5.0), |
| ("priority_weight", 1, 10), |
| ("max_active_tis_per_dag", 1, 10), |
| ("on_execute_callback", [], [id]), |
| ("on_failure_callback", [], [id]), |
| ("on_retry_callback", [], [id]), |
| ("on_success_callback", [], [id]), |
| ("on_skipped_callback", [], [id]), |
| ("inlets", ["a"], ["b"]), |
| ("outlets", ["a"], ["b"]), |
| ("render_template_as_native_obj", True, False), |
| ], |
| ) |
| def test_setters(setter_name: str, old_value: object, new_value: object) -> None: |
| op = MockOperator.partial(task_id="a", arg1="a").expand(arg2=["a", "b", "c"]) |
| setattr(op, setter_name, old_value) |
| assert getattr(op, setter_name) == old_value |
| setattr(op, setter_name, new_value) |
| assert getattr(op, setter_name) == new_value |
| |
| |
| def test_mapped_operator_in_task_group_no_duplicate_prefix(): |
| """Test that task_id doesn't get duplicated prefix when unmapping a mapped operator in a task group.""" |
| from airflow.sdk.definitions.taskgroup import TaskGroup |
| |
| with DAG("test-dag"): |
| with TaskGroup(group_id="tg1") as tg1: |
| # Create a mapped task within the task group |
| mapped_task = MockOperator.partial(task_id="mapped_task", arg1="a").expand(arg2=["a", "b", "c"]) |
| |
| # Check the mapped operator has correct task_id |
| assert mapped_task.task_id == "tg1.mapped_task" |
| assert mapped_task.task_group == tg1 |
| assert mapped_task.task_group.group_id == "tg1" |
| |
| # Simulate what happens during execution - unmap the operator |
| # unmap expects resolved kwargs |
| unmapped = mapped_task.unmap({"arg2": "a"}) |
| |
| # The unmapped operator should have the same task_id, not a duplicate prefix |
| assert unmapped.task_id == "tg1.mapped_task", f"Expected 'tg1.mapped_task' but got '{unmapped.task_id}'" |
