diff --git a/airflow-core/src/airflow/serialization/definitions/mappedoperator.py b/airflow-core/src/airflow/serialization/definitions/mappedoperator.py index 288f1673136b8..8ec4065c6787c 100644 --- a/airflow-core/src/airflow/serialization/definitions/mappedoperator.py +++ b/airflow-core/src/airflow/serialization/definitions/mappedoperator.py @@ -223,6 +223,10 @@ def ignore_first_depends_on_past(self) -> bool: def wait_for_downstream(self) -> bool: return self._get_partial_kwargs_or_operator_default("wait_for_downstream") + @property + def wait_for_past_depends_before_skipping(self) -> bool: + return self._get_partial_kwargs_or_operator_default("wait_for_past_depends_before_skipping") + @property def retries(self) -> int: return self._get_partial_kwargs_or_operator_default("retries") diff --git a/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py b/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py index 80c3380f87745..41abe26d97200 100644 --- a/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py +++ b/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py @@ -140,7 +140,9 @@ def _skip( # If the parent SkipMixin has run, and the XCom result stored indicates this # ti should be skipped, set ti.state to SKIPPED and fail the rule so that the # ti does not execute. - if dep_context.wait_for_past_depends_before_skipping: + if TYPE_CHECKING: + assert ti.task + if dep_context.wait_for_past_depends_before_skipping or ti.task.wait_for_past_depends_before_skipping: if not is_past_depends_met(ti, session=session): yield self._failing_status(reason="Task should be skipped but the past depends are not met") return diff --git a/airflow-core/src/airflow/ti_deps/deps/prev_dagrun_dep.py b/airflow-core/src/airflow/ti_deps/deps/prev_dagrun_dep.py index 1762abb49cd1f..6814fd6616d38 100644 --- a/airflow-core/src/airflow/ti_deps/deps/prev_dagrun_dep.py +++ b/airflow-core/src/airflow/ti_deps/deps/prev_dagrun_dep.py @@ -64,9 +64,12 @@ class PrevDagrunDep(BaseTIDep): IS_TASK_DEP = True @staticmethod - def _push_past_deps_met_xcom_if_needed(ti: TI, dep_context): - if dep_context.wait_for_past_depends_before_skipping: - ti.xcom_push(key=PAST_DEPENDS_MET, value=True) + def _push_past_deps_met_xcom_if_needed(ti: TI, dep_context, *, session: Session): + if TYPE_CHECKING: + assert ti.task + if dep_context.wait_for_past_depends_before_skipping or ti.task.wait_for_past_depends_before_skipping: + # The scheduler's session: a fresh one would commit inside the scheduling loop. + ti.xcom_push(key=PAST_DEPENDS_MET, value=True, session=session) @staticmethod def _has_tis(dagrun: DagRun, task_id: str, *, session: Session) -> bool: @@ -148,13 +151,13 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session): if TYPE_CHECKING: assert ti.task if dep_context.ignore_depends_on_past: - self._push_past_deps_met_xcom_if_needed(ti, dep_context) + self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session) reason = "The context specified that the state of past DAGs could be ignored." yield self._passing_status(reason=reason) return if not ti.task.depends_on_past: - self._push_past_deps_met_xcom_if_needed(ti, dep_context) + self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session) yield self._passing_status(reason="The task did not have depends_on_past set.") return @@ -171,7 +174,7 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session): return if not dr: - self._push_past_deps_met_xcom_if_needed(ti, dep_context) + self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session) yield self._passing_status(reason="This task instance does not belong to a DAG.") return @@ -183,7 +186,7 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session): last_dagrun = DagRun.get_previous_dagrun(dr, session=session) # First ever run for this DAG. if not last_dagrun: - self._push_past_deps_met_xcom_if_needed(ti, dep_context) + self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session) yield self._passing_status(reason="This task instance was the first task instance for its task.") return @@ -194,14 +197,14 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session): and last_dagrun.logical_date is not None and last_dagrun.logical_date < ti.task.start_date ): - self._push_past_deps_met_xcom_if_needed(ti, dep_context) + self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session) yield self._passing_status(reason="This task instance was the first task instance for its task.") return if not self._has_tis(last_dagrun, ti.task_id, session=session): if ti.task.ignore_first_depends_on_past: if not self._has_any_prior_tis(ti, session=session): - self._push_past_deps_met_xcom_if_needed(ti, dep_context) + self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session) yield self._passing_status( reason="ignore_first_depends_on_past is true for this task " "and it is the first task instance for its task." @@ -234,4 +237,4 @@ def _get_dep_statuses(self, ti: TI, dep_context, *, session: Session): ) return - self._push_past_deps_met_xcom_if_needed(ti, dep_context) + self._push_past_deps_met_xcom_if_needed(ti, dep_context, session=session) diff --git a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py index 9ccc4adcb133b..bcd5c3e2cb9b5 100644 --- a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py +++ b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py @@ -366,9 +366,9 @@ def _evaluate_setup_constraint( new_state = TaskInstanceState.REMOVED if new_state is not None: - if ( - new_state == TaskInstanceState.SKIPPED - and dep_context.wait_for_past_depends_before_skipping + if new_state == TaskInstanceState.SKIPPED and ( + dep_context.wait_for_past_depends_before_skipping + or task.wait_for_past_depends_before_skipping ): if not is_past_depends_met(ti, session=session): yield ( @@ -524,9 +524,9 @@ def _evaluate_direct_relatives() -> Iterator[TIDepStatus]: # 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 + if new_state == TaskInstanceState.SKIPPED and ( + dep_context.wait_for_past_depends_before_skipping + or task.wait_for_past_depends_before_skipping ): if not is_past_depends_met(ti, session=session): yield self._failing_status( diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 5c191505666ca..3ab0958366809 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -81,10 +81,14 @@ from airflow.models.taskreschedule import TaskReschedule from airflow.models.trigger import Trigger from airflow.models.variable import Variable -from airflow.models.xcom import XComModelV2 +from airflow.models.xcom import XComModel, XComModelV2 from airflow.providers.standard.operators.bash import BashOperator from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.providers.standard.operators.python import PythonOperator, ShortCircuitOperator +from airflow.providers.standard.operators.python import ( + BranchPythonOperator, + PythonOperator, + ShortCircuitOperator, +) from airflow.sdk import ( DAG, BaseOperator, @@ -103,6 +107,11 @@ from airflow.settings import get_policy_plugin_manager from airflow.task.trigger_rule import TriggerRule from airflow.ti_deps.dep_context import DepContext +from airflow.ti_deps.deps.not_previously_skipped_dep import ( + XCOM_SKIPMIXIN_FOLLOWED, + XCOM_SKIPMIXIN_KEY, + XCOM_SKIPMIXIN_SKIPPED, +) from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep from airflow.triggers.base import StartTriggerArgs from airflow.utils.session import create_session @@ -1284,6 +1293,74 @@ def test_depends_on_past(self, dag_maker, session, prev_ti_state, is_ti_schedula schedulable_tis = [ti.task_id for ti in decision.schedulable_tis] assert ("test_dop_task" in schedulable_tis) == is_ti_schedulable + @pytest.mark.parametrize( + ("skipmixin_operator", "decision"), + [ + pytest.param(ShortCircuitOperator, {XCOM_SKIPMIXIN_SKIPPED: ["past"]}, id="short_circuit"), + pytest.param(BranchPythonOperator, {XCOM_SKIPMIXIN_FOLLOWED: []}, id="branch"), + ], + ) + @pytest.mark.parametrize( + ("wait", "state_while_past_failed"), + [ + pytest.param(False, TaskInstanceState.SKIPPED, id="skip_now"), + pytest.param(True, None, id="wait_for_past_depends"), + ], + ) + def test_skip_decision_waits_for_past_depends( + self, dag_maker, session, skipmixin_operator, decision, wait, state_while_past_failed + ): + """A task told to skip waits for its past depends when it asks to, and is skipped once they are met.""" + with dag_maker( + dag_id="test_skip_decision_waits_for_past_depends", + schedule=datetime.timedelta(days=1), + session=session, + serialized=True, + ): + gate = skipmixin_operator(task_id="gate", python_callable=lambda: None) + past = BaseOperator( + task_id="past", depends_on_past=True, wait_for_past_depends_before_skipping=wait + ) + gate >> past + + dag_run_1: DagRun = dag_maker.create_dagrun( + logical_date=timezone.datetime(2016, 1, 1, 0, 0, 0), run_type=DagRunType.SCHEDULED + ) + dag_run_2: DagRun = dag_maker.create_dagrun( + logical_date=timezone.datetime(2016, 1, 2, 0, 0, 0), run_type=DagRunType.SCHEDULED + ) + # Run 1: the gate ran and the past-dependent task failed. + dag_run_1.get_task_instance("gate", session=session).set_state( + TaskInstanceState.SUCCESS, session=session + ) + past_ti_1 = dag_run_1.get_task_instance("past", session=session) + past_ti_1.set_state(TaskInstanceState.FAILED, session=session) + # Run 2: the gate decided that "past" should be skipped. + dag_run_2.get_task_instance("gate", session=session).set_state( + TaskInstanceState.SUCCESS, session=session + ) + XComModel.set( + key=XCOM_SKIPMIXIN_KEY, + value=decision, + dag_id=dag_run_2.dag_id, + task_id="gate", + run_id=dag_run_2.run_id, + session=session, + ) + session.flush() + + dag_run_2.task_instance_scheduling_decisions(session=session) + session.flush() + assert dag_run_2.get_task_instance("past", session=session).state == state_while_past_failed + + past_ti_1.set_state(TaskInstanceState.SUCCESS, session=session) + session.flush() + # The past-depends-met mark and the skip may land in separate scheduling passes. + for _ in range(2): + dag_run_2.task_instance_scheduling_decisions(session=session) + session.flush() + assert dag_run_2.get_task_instance("past", session=session).state == TaskInstanceState.SKIPPED + @pytest.mark.parametrize( ("prev_ti_state", "is_ti_schedulable"), [ diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index d40f36ebf6d33..9f3d983e58856 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -2075,6 +2075,18 @@ def test_serialize_mapped_outlets(self): assert round_tripped.inlets == [] assert round_tripped.outlets == [] + def test_serialize_mapped_wait_for_past_depends_before_skipping(self): + with DAG(dag_id="d", schedule=None, start_date=datetime.now()): + op = MockOperator.partial( + task_id="x", depends_on_past=True, wait_for_past_depends_before_skipping=True + ).expand(arg1=[1, 2]) + + serialized = OperatorSerialization.serialize_mapped_operator(op) + round_tripped = OperatorSerialization.deserialize_operator(serialized) + assert isinstance(round_tripped, SerializedMappedOperator) + assert round_tripped.depends_on_past is True + assert round_tripped.wait_for_past_depends_before_skipping is True + @pytest.mark.db_test @pytest.mark.parametrize("mapped", [False, True]) def test_derived_dag_deps_sensor(self, mapped): diff --git a/airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py index 0c0fa53b5245f..9d1a97791af1a 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py @@ -152,6 +152,44 @@ def test_parent_skip_branch(session, dag_maker): assert tis["op2"].state == State.SKIPPED +def test_parent_skip_branch_waits_for_past_depends_when_the_task_asks(session, dag_maker): + """ + A BranchPythonOperator does not follow op2, which waits for its past depends before being skipped. + NotPreviouslySkippedDep is not met and leaves op2 unchanged until the past depends are met. + """ + start_date = pendulum.datetime(2020, 1, 1) + with dag_maker( + "test_parent_skip_branch_waits_for_past_depends_dag", + schedule=None, + start_date=start_date, + session=session, + ): + op1 = BranchPythonOperator(task_id="op1", python_callable=lambda: "op3") + op2 = EmptyOperator(task_id="op2", depends_on_past=True, wait_for_past_depends_before_skipping=True) + op3 = EmptyOperator(task_id="op3") + op1 >> [op2, op3] + + dagrun = dag_maker.create_dagrun(run_type=DagRunType.MANUAL, state=State.RUNNING) + tis = {ti.task_id: ti for ti in dagrun.task_instances} + run_task_instance(tis["op1"], op1) + + dep = NotPreviouslySkippedDep() + (status,) = dep.get_dep_statuses(tis["op2"], DepContext(), session=session) + assert "past depends are not met" in status.reason + assert tis["op2"].state == State.NONE + + XComModel.set( + key=PAST_DEPENDS_MET, + value=True, + dag_id=dagrun.dag_id, + task_id="op2", + run_id=dagrun.run_id, + session=session, + ) + assert not dep.is_met(tis["op2"], session=session) + assert tis["op2"].state == State.SKIPPED + + def test_parent_not_executed(session, dag_maker): """ A simple DAG with a BranchPythonOperator that does not follow op2. Parent task is not yet diff --git a/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py index 59a2274c5feb8..1c7a6f601eaae 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py @@ -330,6 +330,20 @@ def test_dagrun_dep(mock_get_previous_dagrun, mock_get_previous_scheduled_dagrun assert actual_dep_met == expected_dep_met if past_depends_met_xcom_sent: - ti.xcom_push.assert_called_with(key="past_depends_met", value=True) + ti.xcom_push.assert_called_with(key="past_depends_met", value=True, session=ANY) else: ti.xcom_push.assert_not_called() + + +def test_past_depends_met_xcom_sent_when_the_task_asks_to_wait(): + """The task's own wait_for_past_depends_before_skipping sends the mark like the context flag does.""" + task = BaseOperator( + task_id="test_task", + dag=DAG("test_dag", schedule=timedelta(days=1), start_date=datetime(2016, 1, 1)), + start_date=datetime(2016, 1, 1), + wait_for_past_depends_before_skipping=True, + ) + ti = Mock(task=task, task_id=task.task_id, **{"xcom_push.return_value": None}) + + assert PrevDagrunDep().is_met(ti=ti, dep_context=DepContext(ignore_depends_on_past=True)) + ti.xcom_push.assert_called_once_with(key="past_depends_met", value=True, session=ANY) diff --git a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py index 27fb20d1b785b..5369e5b904740 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py @@ -734,6 +734,40 @@ def test_all_success_tr_skip_wait_for_past_depends_before_skipping_past_depends_ expected_reason="requires all upstream tasks to have succeeded, but found 1", ) + @pytest.mark.parametrize( + ("past_depends_met", "expected_ti_state", "expected_reason"), + [ + pytest.param(None, None, "Task should be skipped but the past depends are not met", id="unmet"), + pytest.param( + True, SKIPPED, "requires all upstream tasks to have succeeded, but found 1", id="met" + ), + ], + ) + def test_all_success_tr_skip_waits_when_the_task_asks( + self, session, get_task_instance, past_depends_met, expected_ti_state, expected_reason + ): + """The task's own wait_for_past_depends_before_skipping holds the skip like the context flag does.""" + ti = get_task_instance( + TriggerRule.ALL_SUCCESS, + success=1, + skipped=1, + failed=0, + removed=0, + upstream_failed=0, + done=2, + normal_tasks=["FakeTaskID", "OtherFakeTaskID"], + ) + ti.task.wait_for_past_depends_before_skipping = True + xcom_mock = Mock(return_value=past_depends_met) + with mock.patch("airflow.models.taskinstance.TaskInstance.xcom_pull", xcom_mock): + _test_trigger_rule( + ti=ti, + session=session, + flag_upstream_failed=True, + expected_ti_state=expected_ti_state, + expected_reason=expected_reason, + ) + @pytest.mark.parametrize("flag_upstream_failed", [True, False]) def test_none_failed_tr_success(self, session, get_task_instance, flag_upstream_failed): """ @@ -2355,6 +2389,38 @@ def test_wait_for_past_depends_before_skipping_reads_stored_value( m.assert_not_called() +@pytest.mark.parametrize( + ("past_depends_met", "expected_ti_state", "expect_failure"), + [pytest.param(None, None, True, id="unmet"), pytest.param(True, SKIPPED, False, id="met")], +) +def test_setup_constraint_waits_when_the_task_asks( + session, get_task_instance, past_depends_met, expected_ti_state, expect_failure +): + """Setup task with a skipped upstream task: the task's own flag holds the skip like the context flag does.""" + ti = get_task_instance( + trigger_rule=TriggerRule.ALL_DONE, + success=1, + skipped=1, + failed=0, + removed=0, + upstream_failed=0, + done=2, + setup_tasks=["FakeTaskID", "OtherFakeTaskID"], + ) + ti.task.wait_for_past_depends_before_skipping = True + xcom_mock = Mock(return_value=past_depends_met) + with mock.patch("airflow.models.taskinstance.TaskInstance.xcom_pull", xcom_mock): + _test_trigger_rule( + ti=ti, + session=session, + flag_upstream_failed=True, + expected_ti_state=expected_ti_state, + expected_reason=( + "Task should be skipped but the past depends are not met" if expect_failure else "" + ), + ) + + @pytest.mark.parametrize(("flag_upstream_failed", "expected_ti_state"), [(True, SKIPPED), (False, None)]) @pytest.mark.need_serialized_dag def test_setup_mapped_task_group_finished_upstream_before_expand( diff --git a/task-sdk/src/airflow/sdk/bases/skipmixin.py b/task-sdk/src/airflow/sdk/bases/skipmixin.py index f4dd25f5f84ab..f759635136059 100644 --- a/task-sdk/src/airflow/sdk/bases/skipmixin.py +++ b/task-sdk/src/airflow/sdk/bases/skipmixin.py @@ -55,6 +55,24 @@ def _set_state_to_skipped( if tasks and map_index == -1: raise DownstreamTasksSkipped(tasks=tasks) + def _tasks_to_skip_now(self, tasks: Sequence[Operator]) -> list[Operator]: + """ + Drop the direct downstream tasks that wait for their past depends before being skipped. + + They stay in the XCom decision, so the scheduler skips them once the previous run's task + instance has succeeded or been skipped (``NotPreviouslySkippedDep``). + """ + direct_downstream_ids: set[str] = getattr(self, "downstream_task_ids", set()) + return [ + t + for t in tasks + if not ( + t.task_id in direct_downstream_ids + and t.depends_on_past + and t.wait_for_past_depends_before_skipping + ) + ] + def skip( self, ti: RuntimeTaskInstanceProtocol, @@ -79,7 +97,7 @@ def skip( value={XCOM_SKIPMIXIN_SKIPPED: task_ids_list}, ) - self._set_state_to_skipped(task_ids_list, ti.map_index) + self._set_state_to_skipped([t.task_id for t in self._tasks_to_skip_now(task_list)], ti.map_index) def skip_all_except( self, @@ -156,7 +174,9 @@ def skip_all_except( branch_task_id_set.update(dag.get_task(branch_task_id).get_flat_relative_ids(upstream=False)) skip_tasks = [ - (t.task_id, ti.map_index) for t in downstream_tasks if t.task_id not in branch_task_id_set + (t.task_id, ti.map_index) + for t in self._tasks_to_skip_now(downstream_tasks) + if t.task_id not in branch_task_id_set ] follow_task_ids = [t.task_id for t in downstream_tasks if t.task_id in branch_task_id_set] diff --git a/task-sdk/tests/task_sdk/bases/test_skipmixin.py b/task-sdk/tests/task_sdk/bases/test_skipmixin.py index 9518270e2294a..ebdb747159663 100644 --- a/task-sdk/tests/task_sdk/bases/test_skipmixin.py +++ b/task-sdk/tests/task_sdk/bases/test_skipmixin.py @@ -96,6 +96,87 @@ def test_skip_without_task_id_does_not_push_xcom(self): ti.xcom_push.assert_not_called() + def test_skip_leaves_task_waiting_for_past_depends_to_the_scheduler(self): + """skip() leaves a direct downstream task that waits for its past depends out of the state change.""" + mixin = SkipMixin() + mixin.task_id = "gate" + mixin.downstream_task_ids = {"skip_now", "flag_only", "wait"} + + skip_now = MagicMock( + spec=BaseOperator, + task_id="skip_now", + depends_on_past=True, + wait_for_past_depends_before_skipping=False, + ) + flag_only = MagicMock( + spec=BaseOperator, + task_id="flag_only", + depends_on_past=False, + wait_for_past_depends_before_skipping=True, + ) + wait = MagicMock( + spec=BaseOperator, + task_id="wait", + depends_on_past=True, + wait_for_past_depends_before_skipping=True, + ) + # Reached through "wait", not a direct downstream task: skipped at once as before. + indirect = MagicMock( + spec=BaseOperator, + task_id="indirect", + depends_on_past=True, + wait_for_past_depends_before_skipping=True, + ) + ti = Mock(spec=RuntimeTaskInstanceProtocol, map_index=-1) + + with pytest.raises(DownstreamTasksSkipped) as exc_info: + mixin.skip(ti=ti, tasks=[skip_now, flag_only, wait, indirect]) + + # The decision still lists "wait", so NotPreviouslySkippedDep skips it once its past depends are met. + ti.xcom_push.assert_called_once_with( + key=XCOM_SKIPMIXIN_KEY, + value={XCOM_SKIPMIXIN_SKIPPED: ["skip_now", "flag_only", "wait", "indirect"]}, + ) + assert exc_info.value.tasks == ["skip_now", "flag_only", "indirect"] + + def test_skip_all_except_leaves_task_waiting_for_past_depends_to_the_scheduler(self): + """skip_all_except() leaves a task that waits for its past depends out of the state change.""" + mixin = SkipMixin() + mixin.downstream_task_ids = {"down1", "down2", "down3"} + + downstream1 = MagicMock( + spec=BaseOperator, + task_id="down1", + depends_on_past=True, + wait_for_past_depends_before_skipping=True, + ) + downstream2 = MagicMock( + spec=BaseOperator, + task_id="down2", + depends_on_past=True, + wait_for_past_depends_before_skipping=False, + ) + downstream3 = MagicMock(spec=BaseOperator, task_id="down3") + downstream3.get_flat_relative_ids.return_value = set() + + mock_task = MagicMock(spec=BaseOperator) + mock_task.downstream_list = [downstream1, downstream2, downstream3] + mock_dag = MagicMock(spec=DAG) + mock_dag.task_ids = ["gate", "down1", "down2", "down3"] + mock_dag.get_task.return_value = downstream3 + mock_task.dag = mock_dag + + ti = Mock(spec=RuntimeTaskInstanceProtocol, map_index=-1, task=mock_task) + + with pytest.raises(DownstreamTasksSkipped) as exc_info: + mixin.skip_all_except(ti=ti, branch_task_ids="down3") + + assert exc_info.value.tasks == [("down2", -1)] + ti.xcom_push.assert_called_once_with( + key=XCOM_SKIPMIXIN_KEY, + value={XCOM_SKIPMIXIN_FOLLOWED: ["down3"]}, + ) + def test_skip_all_except_with_none_skips_all(self): """skip_all_except(None) should skip all downstream tasks.""" mixin = SkipMixin()