diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py index ba357c0bd1bbf..81e6cd48eec27 100644 --- a/airflow/models/abstractoperator.py +++ b/airflow/models/abstractoperator.py @@ -293,6 +293,17 @@ def get_upstreams_only_setups_and_teardowns(self) -> Iterable[Operator]: if t.is_teardown and not t == self: yield t + def get_upstreams_only_setups(self) -> Iterable[Operator]: + """ + Only upstream setups. + + This method is meant to be used when we are checking task dependencies where we need + to wait for all the upstream setups to complete before we can run the task. + """ + for task in self.get_upstreams_only_setups_and_teardowns(): + if task.is_setup: + yield task + def _iter_all_mapped_downstreams(self) -> Iterator[MappedOperator | MappedTaskGroup]: """Return mapped nodes that are direct dependencies of the current task. diff --git a/airflow/models/dag.py b/airflow/models/dag.py index cc2626278f2ee..db6723a97ef84 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -128,6 +128,7 @@ with_row_locks, ) from airflow.utils.state import DagRunState, TaskInstanceState +from airflow.utils.trigger_rule import TriggerRule from airflow.utils.types import NOTSET, ArgNotSet, DagRunType, EdgeInfoType if TYPE_CHECKING: @@ -717,6 +718,11 @@ def validate_setup_teardown(self): :meta private: """ for task in self.tasks: + if task.is_setup: + for down_task in task.downstream_list: + if not down_task.is_teardown and down_task.trigger_rule != TriggerRule.ALL_SUCCESS: + # this is required to ensure consistent clearing behavior when upstream + raise ValueError("Setup tasks must be followed with trigger rule ALL_SUCCESS.") FailStopDagInvalidTriggerRule.check(dag=self, trigger_rule=task.trigger_rule) def __repr__(self): diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py index dbdf692e769ec..d73b6dd176694 100644 --- a/airflow/ti_deps/deps/trigger_rule_dep.py +++ b/airflow/ti_deps/deps/trigger_rule_dep.py @@ -20,10 +20,12 @@ import collections import collections.abc import functools -from typing import TYPE_CHECKING, Iterator, NamedTuple +from dataclasses import dataclass +from typing import TYPE_CHECKING, Iterator from sqlalchemy import and_, func, or_, select +from airflow.models import MappedOperator from airflow.models.taskinstance import PAST_DEPENDS_MET from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.deps.base_ti_dep import BaseTIDep, TIDepStatus @@ -34,10 +36,13 @@ from sqlalchemy.orm import Session from sqlalchemy.sql.expression import ColumnOperators + from airflow.models.dagrun import DagRun + from airflow.models.operator import Operator from airflow.models.taskinstance import TaskInstance -class _UpstreamTIStates(NamedTuple): +@dataclass +class _UpstreamTIStates: """States of the upstream tis for a specific ti. This is used to determine whether the specific ti can run in this iteration. @@ -51,9 +56,12 @@ class _UpstreamTIStates(NamedTuple): done: int success_setup: int skipped_setup: int + failed_setup: int = 0 @classmethod - def calculate(cls, finished_upstreams: Iterator[TaskInstance]) -> _UpstreamTIStates: + def calculate( + cls, finished_upstreams: Iterator[TaskInstance], finished_setup_upstream_tis + ) -> _UpstreamTIStates: """Calculate states for a task instance. ``counter`` is inclusive of ``setup_counter`` -- e.g. if there are 2 skipped upstreams, one @@ -67,8 +75,9 @@ def calculate(cls, finished_upstreams: Iterator[TaskInstance]) -> _UpstreamTISta for ti in finished_upstreams: curr_state = {ti.state: 1} counter.update(curr_state) - if ti.task.is_setup: - setup_counter.update(curr_state) + for ti in finished_setup_upstream_tis: + curr_state = {ti.state: 1} + setup_counter.update(curr_state) return _UpstreamTIStates( success=counter.get(TaskInstanceState.SUCCESS, 0), skipped=counter.get(TaskInstanceState.SKIPPED, 0), @@ -78,6 +87,8 @@ def calculate(cls, finished_upstreams: Iterator[TaskInstance]) -> _UpstreamTISta done=sum(counter.values()), success_setup=setup_counter.get(TaskInstanceState.SUCCESS, 0), skipped_setup=setup_counter.get(TaskInstanceState.SKIPPED, 0), + failed_setup=setup_counter.get(TaskInstanceState.FAILED, 0) + + setup_counter.get(TaskInstanceState.UPSTREAM_FAILED, 0), ) @@ -121,6 +132,13 @@ def _evaluate_trigger_rule( from airflow.models.operator import needs_expansion from airflow.models.taskinstance import TaskInstance + if ti.task.is_teardown: + setup_upstream_tasks = [task for task in ti.task.upstream_list if task.is_setup] + else: + setup_upstream_tasks = list(ti.task.get_upstreams_only_setups()) + + setup_upstream_tasks_ids = {task.task_id for task in setup_upstream_tasks} + task = ti.task upstream_tasks = {t.task_id: t for t in task.upstream_list} trigger_rule = task.trigger_rule @@ -152,6 +170,26 @@ def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None: session=session, ) + def _get_tis_count_from_tasks_list(dag_run: DagRun, tasks_ids_list: list[str]) -> int: + """Get the number of task instances from a list of tasks.""" + return session.execute( + select(func.count("*")) + .where(TaskInstance.dag_id == dag_run.dag_id, TaskInstance.run_id == dag_run.run_id) + .where(TaskInstance.task_id.in_(tasks_ids_list)) + ).scalar_one() + + def _get_setup_upstream_tis_count( + current_ti: TaskInstance, setup_upstream_tasks_list: list[Operator] + ): + mapped_setup_upstream_ids = { + task.task_id + for task in setup_upstream_tasks_list + if isinstance(task, MappedOperator) or task.task_group != current_ti.task.task_group + } + return _get_tis_count_from_tasks_list(current_ti.dag_run, list(mapped_setup_upstream_ids)) + len( + [task for task in setup_upstream_tasks_list if task.task_id not in mapped_setup_upstream_ids] + ) + def _is_relevant_upstream(upstream: TaskInstance) -> bool: """Whether a task instance is a "relevant upstream" of the current task.""" # Not actually an upstream task. @@ -176,12 +214,32 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool: return True return False + def _is_relevant_setup_upstream(upstream: TaskInstance) -> bool: + """Whether a task instance is a "relevant upstream setup" of the current task.""" + if upstream.task_id in setup_upstream_tasks_ids: + relevant = _get_relevant_upstream_map_indexes(upstream.task_id) + if relevant is None: + return True + if relevant == upstream.map_index: + return True + if isinstance(relevant, collections.abc.Container) and upstream.map_index in relevant: + return True + return False + return False + finished_upstream_tis = ( finished_ti for finished_ti in dep_context.ensure_finished_tis(ti.get_dagrun(session), session) if _is_relevant_upstream(finished_ti) ) - upstream_states = _UpstreamTIStates.calculate(finished_upstream_tis) + + finished_setup_upstream_tis = ( + finished_ti + for finished_ti in dep_context.ensure_finished_tis(ti.get_dagrun(session), session) + if _is_relevant_setup_upstream(finished_ti) + ) + + upstream_states = _UpstreamTIStates.calculate(finished_upstream_tis, finished_setup_upstream_tis) success = upstream_states.success skipped = upstream_states.skipped @@ -191,6 +249,7 @@ def _is_relevant_upstream(upstream: TaskInstance) -> bool: done = upstream_states.done success_setup = upstream_states.success_setup skipped_setup = upstream_states.skipped_setup + failed_setup = upstream_states.failed_setup def _iter_upstream_conditions() -> Iterator[ColumnOperators]: # Optimization: If the current task is not in a mapped task group, @@ -223,9 +282,12 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]: # Optimization: Don't need to hit the database if all upstreams are # "simple" tasks (no task or task group mapping involved). + if not ti.task.is_teardown: + upstream_setup = _get_setup_upstream_tis_count(ti, setup_upstream_tasks) if not any(needs_expansion(t) for t in upstream_tasks.values()): upstream = len(upstream_tasks) - upstream_setup = sum(1 for x in upstream_tasks.values() if x.is_setup) + if ti.task.is_teardown: + upstream_setup = sum(1 for x in upstream_tasks.values() if x.is_setup) else: task_id_counts = session.execute( select(TaskInstance.task_id, func.count(TaskInstance.task_id)) @@ -234,59 +296,72 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]: .group_by(TaskInstance.task_id) ).all() upstream = sum(count for _, count in task_id_counts) - upstream_setup = sum(c for t, c in task_id_counts if upstream_tasks[t].is_setup) + if ti.task.is_teardown: + upstream_setup = sum(c for t, c in task_id_counts if upstream_tasks[t].is_setup) upstream_done = done >= upstream + setup_done = (success_setup + skipped_setup + failed_setup) >= upstream_setup changed = False new_state = None if dep_context.flag_upstream_failed: - if trigger_rule == TR.ALL_SUCCESS: - if upstream_failed or failed: - new_state = TaskInstanceState.UPSTREAM_FAILED - elif skipped: - new_state = TaskInstanceState.SKIPPED - elif removed and success and ti.map_index > -1: - if ti.map_index >= success: - new_state = TaskInstanceState.REMOVED - elif trigger_rule == TR.ALL_FAILED: - if success or skipped: - new_state = TaskInstanceState.SKIPPED - elif trigger_rule == TR.ONE_SUCCESS: - if upstream_done and done == skipped: - # if upstream is done and all are skipped mark as skipped - new_state = TaskInstanceState.SKIPPED - elif upstream_done and success <= 0: - # if upstream is done and there are no success mark as upstream failed - new_state = TaskInstanceState.UPSTREAM_FAILED - elif trigger_rule == TR.ONE_FAILED: - if upstream_done and not (failed or upstream_failed): - new_state = TaskInstanceState.SKIPPED - elif trigger_rule == TR.ONE_DONE: - if upstream_done and not (failed or success): - new_state = TaskInstanceState.SKIPPED - elif trigger_rule == TR.NONE_FAILED: - if upstream_failed or failed: - new_state = TaskInstanceState.UPSTREAM_FAILED - elif trigger_rule == TR.NONE_FAILED_MIN_ONE_SUCCESS: - if upstream_failed or failed: - new_state = TaskInstanceState.UPSTREAM_FAILED - elif skipped == upstream: - new_state = TaskInstanceState.SKIPPED - elif trigger_rule == TR.NONE_SKIPPED: - if skipped: - new_state = TaskInstanceState.SKIPPED - elif trigger_rule == TR.ALL_SKIPPED: - if success or failed: - new_state = TaskInstanceState.SKIPPED - elif trigger_rule == TR.ALL_DONE_SETUP_SUCCESS: - if upstream_done and upstream_setup and skipped_setup >= upstream_setup: - # when there is an upstream setup and they have all skipped, then skip - new_state = TaskInstanceState.SKIPPED - elif upstream_done and upstream_setup and success_setup == 0: - # when there is an upstream setup, if none succeeded, mark upstream failed - # if at least one setup ran, we'll let it run - new_state = TaskInstanceState.UPSTREAM_FAILED + if not task.is_teardown and failed_setup: + # we should exclude the teardown tasks from this check, + # because they should be run even if there is only one success setup task + new_state = TaskInstanceState.UPSTREAM_FAILED + elif not task.is_teardown and upstream_setup and setup_done and skipped_setup > 0: + # when there are upstream setup tasks and at least one of them is skipped, then skip + new_state = TaskInstanceState.SKIPPED + elif not upstream_setup or setup_done: + # if there are no upstream setup tasks or all of them are done, + # and we haven't set a new state, then we can check the upstream tasks + if trigger_rule == TR.ALL_SUCCESS: + if upstream_failed or failed: + new_state = TaskInstanceState.UPSTREAM_FAILED + elif skipped: + new_state = TaskInstanceState.SKIPPED + elif removed and success and ti.map_index > -1: + if ti.map_index >= success: + new_state = TaskInstanceState.REMOVED + elif trigger_rule == TR.ALL_FAILED: + if success or skipped: + new_state = TaskInstanceState.SKIPPED + elif trigger_rule == TR.ONE_SUCCESS: + if upstream_done and done == skipped: + # if upstream is done and all are skipped mark as skipped + new_state = TaskInstanceState.SKIPPED + elif upstream_done and success <= 0: + # if upstream is done and there are no success mark as upstream failed + new_state = TaskInstanceState.UPSTREAM_FAILED + elif trigger_rule == TR.ONE_FAILED: + if upstream_done and not (failed or upstream_failed): + new_state = TaskInstanceState.SKIPPED + elif trigger_rule == TR.ONE_DONE: + if upstream_done and not (failed or success): + new_state = TaskInstanceState.SKIPPED + elif trigger_rule == TR.NONE_FAILED: + if upstream_failed or failed: + new_state = TaskInstanceState.UPSTREAM_FAILED + elif trigger_rule == TR.NONE_FAILED_MIN_ONE_SUCCESS: + if upstream_failed or failed: + new_state = TaskInstanceState.UPSTREAM_FAILED + elif skipped == upstream: + new_state = TaskInstanceState.SKIPPED + elif trigger_rule == TR.NONE_SKIPPED: + if skipped: + new_state = TaskInstanceState.SKIPPED + elif trigger_rule == TR.ALL_SKIPPED: + if success or failed: + new_state = TaskInstanceState.SKIPPED + elif trigger_rule == TR.ALL_DONE_SETUP_SUCCESS: + if upstream_done and upstream_setup and skipped_setup >= upstream_setup: + # when there is an upstream setup and they have all skipped, then skip + new_state = TaskInstanceState.SKIPPED + elif upstream_done and upstream_setup and setup_done and success_setup == 0: + # when there is an upstream setup, if none succeeded, mark upstream failed + # if at least one setup ran, we'll let it run + new_state = TaskInstanceState.UPSTREAM_FAILED + if new_state is not None: if new_state == TaskInstanceState.SKIPPED and dep_context.wait_for_past_depends_before_skipping: past_depends_met = ti.xcom_pull( @@ -302,6 +377,15 @@ def _iter_upstream_conditions() -> Iterator[ColumnOperators]: if changed: dep_context.have_changed_ti_states = True + # first, check if all the setup upstream tasks are done, if not, we can't run this task + # we should exclude the teardown tasks from this check, because they should be run even + # if there is only one success setup task + if not task.is_teardown and (success_setup + skipped_setup) < upstream_setup: + yield self._failing_status( + reason=f"Waiting {upstream_setup - (success_setup + skipped_setup)}" + " setup task(s) to complete." + ) + if trigger_rule == TR.ONE_SUCCESS: if success <= 0: yield self._failing_status( diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index e50917e101074..4b6fcf39d5937 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -1189,7 +1189,7 @@ def test_depends_on_past(self, dag_maker): param( "all_done_setup_success", 1, - _UpstreamTIStates(5, 0, 1, 0, 0, 6, 0, 0), + _UpstreamTIStates(5, 0, 1, 0, 0, 6, 0, 0, 1), True, State.UPSTREAM_FAILED, False, @@ -1198,7 +1198,7 @@ def test_depends_on_past(self, dag_maker): param( "all_done_setup_success", 2, - _UpstreamTIStates(5, 0, 2, 0, 0, 7, 0, 0), + _UpstreamTIStates(5, 0, 2, 0, 0, 7, 0, 0, 2), True, State.UPSTREAM_FAILED, False, @@ -1225,7 +1225,7 @@ def test_depends_on_past(self, dag_maker): param( "all_done_setup_success", 2, - _UpstreamTIStates(5, 1, 1, 0, 0, 7, 0, 1), + _UpstreamTIStates(5, 1, 1, 0, 0, 7, 0, 1, 1), True, State.UPSTREAM_FAILED, False, @@ -1234,7 +1234,7 @@ def test_depends_on_past(self, dag_maker): param( "all_done_setup_success", 2, - _UpstreamTIStates(6, 0, 1, 0, 0, 7, 1, 0), + _UpstreamTIStates(6, 0, 1, 0, 0, 7, 1, 0, 1), True, None, True, @@ -1304,6 +1304,177 @@ def test_check_task_dependencies( task.set_downstream(downstream) for i in range(upstream_setups): task = EmptyOperator(task_id=f"setup_{i}", dag=dag).as_setup() + + @pytest.mark.parametrize( + "trigger_rule, direct_upstream_setups, indirect_upstream_setups, upstream_states," + " flag_upstream_failed, expect_state, expect_passed", + [ + param( + "all_success", + 0, + 2, + _UpstreamTIStates(7, 0, 0, 0, 0, 7, 2, 0, 0), + True, + None, + True, + id="indirect upstream setups - all success", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 0, 0), + True, + None, + False, + id="indirect upstream setups - one not done and one success", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 1, 0), + True, + None, + False, + id="indirect upstream setups - one not done and one skipped", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 1, 0), + True, + TaskInstanceState.SKIPPED, + True, + id="indirect upstream setups - all done: one skipped and one success", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(7, 0, 0, 0, 0, 7, 1, 0, 1), + True, + TaskInstanceState.UPSTREAM_FAILED, + False, + id="indirect upstream setups - all done: one success and one failed", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 1, 1), + True, + TaskInstanceState.UPSTREAM_FAILED, + False, + id="indirect upstream setups - all done: one skipped and one failed", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(7, 0, 0, 0, 0, 7, 0, 0, 1), + True, + TaskInstanceState.UPSTREAM_FAILED, + False, + id="indirect upstream setups - one not done and one failed", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(0, 0, 1, 0, 0, 1, 2, 0, 0), + True, + TaskInstanceState.UPSTREAM_FAILED, + False, + id="indirect upstream setups - all setup success but upstream failed", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(5, 0, 0, 0, 0, 5, 2, 0, 0), + True, + None, + False, + id="indirect upstream setups - all setup success but upstream still running", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(5, 0, 0, 0, 0, 5, 0, 0, 1), + True, + TaskInstanceState.UPSTREAM_FAILED, + False, + id="indirect upstream setups - one setup failed but upstream still running", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(4, 0, 1, 0, 0, 5, 0, 2, 0), + True, + TaskInstanceState.SKIPPED, + False, + id="indirect upstream setups - all setup skipped but upstream failed", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(0, 7, 0, 0, 0, 7, 0, 0, 1), + True, + TaskInstanceState.UPSTREAM_FAILED, + False, + id="indirect upstream setups - one setup failed but upstream skipped", + ), + param( + "all_success", + 0, + 2, + _UpstreamTIStates(4, 0, 1, 0, 0, 5, 0, 1, 0), + True, + None, + False, + id="indirect upstream setups - one setup still running but upstream failed", + ), + ], + ) + def test_check_task_dependencies_indirect_upstream( + self, + monkeypatch, + dag_maker, + trigger_rule: str, + direct_upstream_setups: int, + indirect_upstream_setups: int, + upstream_states: _UpstreamTIStates, + flag_upstream_failed: bool, + expect_state: State, + expect_passed: bool, + ): + monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states) + + # sanity checks + s = upstream_states + assert s.skipped >= s.skipped_setup - indirect_upstream_setups + assert s.success >= s.success_setup - indirect_upstream_setups + assert s.done == s.failed + s.success + s.removed + s.upstream_failed + s.skipped + + with dag_maker() as dag: + downstream = EmptyOperator(task_id="downstream", trigger_rule=trigger_rule) + if trigger_rule == "all_done_setup_success": + downstream.is_teardown = True + for i in range(5): + task = EmptyOperator(task_id=f"work_{i}", dag=dag) + task.set_downstream(downstream) + for i in range(direct_upstream_setups): + task = EmptyOperator(task_id=f"direct_setup_{i}", dag=dag).as_setup() + task.set_downstream(downstream) + for i in range(indirect_upstream_setups): + setup_task = EmptyOperator(task_id=f"indirect_setup_{i}", dag=dag).as_setup() + task = EmptyOperator(task_id=f"indirect_setup_downstream_{i}", dag=dag) + setup_task.set_downstream(task) task.set_downstream(downstream) assert task.start_date is not None run_date = task.start_date + datetime.timedelta(days=5) diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py b/tests/ti_deps/deps/test_trigger_rule_dep.py index faa70b5a4951d..8e3a424050664 100644 --- a/tests/ti_deps/deps/test_trigger_rule_dep.py +++ b/tests/ti_deps/deps/test_trigger_rule_dep.py @@ -52,6 +52,7 @@ def _get_task_instance( done: int = 0, skipped_setup: int = 0, success_setup: int = 0, + failed_setup: int = 0, normal_tasks: list[str] | None = None, setup_tasks: list[str] | None = None, ): @@ -78,6 +79,7 @@ def _get_task_instance( done=done, skipped_setup=skipped_setup, success_setup=success_setup, + failed_setup=failed_setup, ) monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: fake_upstream_states) @@ -694,14 +696,14 @@ def test_all_done_tr_success(self, session, get_task_instance): ), param( dict(work=2, setup=1), - dict(success=2, done=3), + dict(success=2, done=3, failed_setup=1), "requires at least one upstream setup task be successful", UPSTREAM_FAILED, id="setup failed", ), param( dict(work=2, setup=2), - dict(success=2, done=4, success_setup=1), + dict(success=2, done=4, success_setup=1, failed_setup=1), None, None, id="one setup failed one success", @@ -715,14 +717,14 @@ def test_all_done_tr_success(self, session, get_task_instance): ), param( dict(work=2, setup=1), - dict(success=2, done=3, failed=1), + dict(success=2, done=3, failed=1, failed_setup=1), "requires at least one upstream setup task be successful", UPSTREAM_FAILED, id="setup failed", ), param( dict(work=2, setup=2), - dict(success=2, done=4, failed=1, skipped_setup=1), + dict(success=2, done=4, failed=1, skipped_setup=1, failed_setup=1), "requires at least one upstream setup task be successful", UPSTREAM_FAILED, id="one setup failed one skipped", @@ -763,12 +765,13 @@ def test_teardown_tr_not_all_done( """ All-done trigger rule success """ - ti = get_task_instance( + ti: TaskInstance = get_task_instance( TriggerRule.ALL_DONE_SETUP_SUCCESS, **states, normal_tasks=[f"w{x}" for x in range(task_cfg["work"])], setup_tasks=[f"s{x}" for x in range(task_cfg["setup"])], ) + ti.task.is_teardown = True dep_statuses = tuple( TriggerRuleDep()._evaluate_trigger_rule( ti=ti, dep_context=DepContext(flag_upstream_failed=True), session=session @@ -984,9 +987,15 @@ def _get_finished_tis(task_id: str) -> Iterator[TaskInstance]: return (ti for ti in tis.values() if ti.task_id in tis[task_id].task.upstream_task_ids) # check handling with cases that tasks are triggered from backfill with no finished tasks - assert _UpstreamTIStates.calculate(_get_finished_tis("op2")) == (1, 0, 0, 0, 0, 1, 0, 0) - assert _UpstreamTIStates.calculate(_get_finished_tis("op4")) == (1, 0, 1, 0, 0, 2, 0, 0) - assert _UpstreamTIStates.calculate(_get_finished_tis("op5")) == (2, 0, 1, 0, 0, 3, 0, 0) + assert _UpstreamTIStates.calculate(_get_finished_tis("op2"), []) == _UpstreamTIStates( + 1, 0, 0, 0, 0, 1, 0, 0, 0 + ) + assert _UpstreamTIStates.calculate(_get_finished_tis("op4"), []) == _UpstreamTIStates( + 1, 0, 1, 0, 0, 2, 0, 0, 0 + ) + assert _UpstreamTIStates.calculate(_get_finished_tis("op5"), []) == _UpstreamTIStates( + 2, 0, 1, 0, 0, 3, 0, 0, 0 + ) dr.update_state(session=session) assert dr.state == DagRunState.SUCCESS @@ -1015,6 +1024,7 @@ def test_mapped_task_upstream_removed_with_all_success_trigger_rules( done=5, skipped_setup=0, success_setup=0, + failed_setup=0, ) monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states) @@ -1054,6 +1064,7 @@ def test_mapped_task_upstream_removed_with_all_failed_trigger_rules( done=5, skipped_setup=0, success_setup=0, + failed_setup=0, ) monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states) @@ -1096,6 +1107,7 @@ def test_mapped_task_upstream_removed_with_none_failed_trigger_rules( done=5, skipped_setup=0, success_setup=0, + failed_setup=0, ) monkeypatch.setattr(_UpstreamTIStates, "calculate", lambda *_: upstream_states)