Fix deferred task resumption in `dag.test()` #51182

kaxil · 2025-05-28T22:58:03Z

When using dag.test() with deferred tasks, tasks that complete their trigger execution were incorrectly being set to SUCCESS state instead of SCHEDULED state. This prevented task resumption!

Dag used to test:

from datetime import datetime, timedelta, timezone
from typing import Any

import pendulum

from airflow.providers.standard.triggers.temporal import DateTimeTrigger
from airflow.sdk import Context, task, BaseOperator, DAG

class DummyOperator(BaseOperator):

    def execute(self, context: Context):
        self.defer(
            trigger=DateTimeTrigger(
                moment=datetime.now(timezone.utc) + timedelta(seconds=2),
            ),
            method_name="execute_complet",
        )

    def execute_complet(self, context: Context, event: Any = None):
        assert event is not None
        return "test"

@task
def dummy_task(param):
    print("DEBUG")
    assert param == "test", "Parameter should be 'test'"

with DAG(
    dag_id="example_debug",
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
) as dag:
    task1 = DummyOperator(task_id="task1")
    task2 = dummy_task(task1.output)
    task1 >> task2

if __name__ == "__main__":
    dag.test()

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

@task

When using `dag.test()` with deferred tasks, tasks that complete their trigger execution were incorrectly being set to `SUCCESS` state instead of `SCHEDULED` state. This prevented task resumption! Dag used to test: ```python from datetime import datetime, timedelta, timezone from typing import Any import pendulum from airflow.providers.standard.triggers.temporal import DateTimeTrigger from airflow.sdk import Context, task, BaseOperator, DAG class DummyOperator(BaseOperator): def execute(self, context: Context): self.defer( trigger=DateTimeTrigger( moment=datetime.now(timezone.utc) + timedelta(seconds=2), ), method_name="execute_complet", ) def execute_complet(self, context: Context, event: Any = None): assert event is not None return "test" @task def dummy_task(param): print("DEBUG") assert param == "test", "Parameter should be 'test'" with DAG( dag_id="example_debug", start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), ) as dag: task1 = DummyOperator(task_id="task1") task2 = dummy_task(task1.output) task1 >> task2 if __name__ == "__main__": dag.test() ```

task-sdk/src/airflow/sdk/definitions/dag.py

@task

When using `dag.test()` with deferred tasks, tasks that complete their trigger execution were incorrectly being set to `SUCCESS` state instead of `SCHEDULED` state. This prevented task resumption! Dag used to test: ```python from datetime import datetime, timedelta, timezone from typing import Any import pendulum from airflow.providers.standard.triggers.temporal import DateTimeTrigger from airflow.sdk import Context, task, BaseOperator, DAG class DummyOperator(BaseOperator): def execute(self, context: Context): self.defer( trigger=DateTimeTrigger( moment=datetime.now(timezone.utc) + timedelta(seconds=2), ), method_name="execute_complet", ) def execute_complet(self, context: Context, event: Any = None): assert event is not None return "test" @task def dummy_task(param): print("DEBUG") assert param == "test", "Parameter should be 'test'" with DAG( dag_id="example_debug", start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), ) as dag: task1 = DummyOperator(task_id="task1") task2 = dummy_task(task1.output) task1 >> task2 if __name__ == "__main__": dag.test() ``` (cherry picked from commit 1b83f71) Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

github-actions · 2025-05-29T11:19:12Z

Backport successfully created: v3-0-test

Status	Branch	Result
✅	v3-0-test

@task

) When using `dag.test()` with deferred tasks, tasks that complete their trigger execution were incorrectly being set to `SUCCESS` state instead of `SCHEDULED` state. This prevented task resumption! Dag used to test: ```python from datetime import datetime, timedelta, timezone from typing import Any import pendulum from airflow.providers.standard.triggers.temporal import DateTimeTrigger from airflow.sdk import Context, task, BaseOperator, DAG class DummyOperator(BaseOperator): def execute(self, context: Context): self.defer( trigger=DateTimeTrigger( moment=datetime.now(timezone.utc) + timedelta(seconds=2), ), method_name="execute_complet", ) def execute_complet(self, context: Context, event: Any = None): assert event is not None return "test" @task def dummy_task(param): print("DEBUG") assert param == "test", "Parameter should be 'test'" with DAG( dag_id="example_debug", start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), ) as dag: task1 = DummyOperator(task_id="task1") task2 = dummy_task(task1.output) task1 >> task2 if __name__ == "__main__": dag.test() ``` (cherry picked from commit 1b83f71) Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

@task

…51199) When using `dag.test()` with deferred tasks, tasks that complete their trigger execution were incorrectly being set to `SUCCESS` state instead of `SCHEDULED` state. This prevented task resumption! Dag used to test: ```python from datetime import datetime, timedelta, timezone from typing import Any import pendulum from airflow.providers.standard.triggers.temporal import DateTimeTrigger from airflow.sdk import Context, task, BaseOperator, DAG class DummyOperator(BaseOperator): def execute(self, context: Context): self.defer( trigger=DateTimeTrigger( moment=datetime.now(timezone.utc) + timedelta(seconds=2), ), method_name="execute_complet", ) def execute_complet(self, context: Context, event: Any = None): assert event is not None return "test" @task def dummy_task(param): print("DEBUG") assert param == "test", "Parameter should be 'test'" with DAG( dag_id="example_debug", start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), ) as dag: task1 = DummyOperator(task_id="task1") task2 = dummy_task(task1.output) task1 >> task2 if __name__ == "__main__": dag.test() ``` (cherry picked from commit 1b83f71) Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

vincbeck · 2025-05-29T14:06:07Z

Thank you!

@task

…51199) When using `dag.test()` with deferred tasks, tasks that complete their trigger execution were incorrectly being set to `SUCCESS` state instead of `SCHEDULED` state. This prevented task resumption! Dag used to test: ```python from datetime import datetime, timedelta, timezone from typing import Any import pendulum from airflow.providers.standard.triggers.temporal import DateTimeTrigger from airflow.sdk import Context, task, BaseOperator, DAG class DummyOperator(BaseOperator): def execute(self, context: Context): self.defer( trigger=DateTimeTrigger( moment=datetime.now(timezone.utc) + timedelta(seconds=2), ), method_name="execute_complet", ) def execute_complet(self, context: Context, event: Any = None): assert event is not None return "test" @task def dummy_task(param): print("DEBUG") assert param == "test", "Parameter should be 'test'" with DAG( dag_id="example_debug", start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), ) as dag: task1 = DummyOperator(task_id="task1") task2 = dummy_task(task1.output) task1 >> task2 if __name__ == "__main__": dag.test() ``` (cherry picked from commit 1b83f71) Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

@task

When using `dag.test()` with deferred tasks, tasks that complete their trigger execution were incorrectly being set to `SUCCESS` state instead of `SCHEDULED` state. This prevented task resumption! Dag used to test: ```python from datetime import datetime, timedelta, timezone from typing import Any import pendulum from airflow.providers.standard.triggers.temporal import DateTimeTrigger from airflow.sdk import Context, task, BaseOperator, DAG class DummyOperator(BaseOperator): def execute(self, context: Context): self.defer( trigger=DateTimeTrigger( moment=datetime.now(timezone.utc) + timedelta(seconds=2), ), method_name="execute_complet", ) def execute_complet(self, context: Context, event: Any = None): assert event is not None return "test" @task def dummy_task(param): print("DEBUG") assert param == "test", "Parameter should be 'test'" with DAG( dag_id="example_debug", start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), ) as dag: task1 = DummyOperator(task_id="task1") task2 = dummy_task(task1.output) task1 >> task2 if __name__ == "__main__": dag.test() ```

@task

When using `dag.test()` with deferred tasks, tasks that complete their trigger execution were incorrectly being set to `SUCCESS` state instead of `SCHEDULED` state. This prevented task resumption! Dag used to test: ```python from datetime import datetime, timedelta, timezone from typing import Any import pendulum from airflow.providers.standard.triggers.temporal import DateTimeTrigger from airflow.sdk import Context, task, BaseOperator, DAG class DummyOperator(BaseOperator): def execute(self, context: Context): self.defer( trigger=DateTimeTrigger( moment=datetime.now(timezone.utc) + timedelta(seconds=2), ), method_name="execute_complet", ) def execute_complet(self, context: Context, event: Any = None): assert event is not None return "test" @task def dummy_task(param): print("DEBUG") assert param == "test", "Parameter should be 'test'" with DAG( dag_id="example_debug", start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), ) as dag: task1 = DummyOperator(task_id="task1") task2 = dummy_task(task1.output) task1 >> task2 if __name__ == "__main__": dag.test() ```

kaxil requested a review from vincbeck May 28, 2025 22:58

kaxil requested review from amoghrajesh and ashb as code owners May 28, 2025 22:58

boring-cyborg bot added area:CLI area:task-sdk labels May 28, 2025

kaxil added backport-to-v3-1-test Mark PR with this label to backport to v3-1-test branch and removed area:CLI area:task-sdk labels May 28, 2025

kaxil added this to the Airflow 3.0.2 milestone May 28, 2025

kaxil mentioned this pull request May 28, 2025

Port dag.test to Task SDK #50300

Merged

3 tasks

amoghrajesh reviewed May 29, 2025

View reviewed changes

task-sdk/src/airflow/sdk/definitions/dag.py Show resolved Hide resolved

amoghrajesh approved these changes May 29, 2025

View reviewed changes

kaxil merged commit 1b83f71 into apache:main May 29, 2025
77 checks passed

kaxil deleted the fix-dag-test-bug branch May 29, 2025 11:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix deferred task resumption in `dag.test()` #51182

Fix deferred task resumption in `dag.test()` #51182

Uh oh!

kaxil commented May 28, 2025

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented May 29, 2025

Uh oh!

vincbeck commented May 29, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Fix deferred task resumption in dag.test() #51182

Fix deferred task resumption in dag.test() #51182

Uh oh!

Conversation

kaxil commented May 28, 2025

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented May 29, 2025

Backport successfully created: v3-0-test

Uh oh!

vincbeck commented May 29, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Fix deferred task resumption in `dag.test()` #51182

Fix deferred task resumption in `dag.test()` #51182