This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch v3-0-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-0-test by this push:
new 778823a7a80 [v3-0-test] Fix deferred task resumption in ``dag.test()``
(#51182) (#51199)
778823a7a80 is described below
commit 778823a7a805dbe6e9393fc4b757220dbb1faaba
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Thu May 29 19:31:11 2025 +0530
[v3-0-test] Fix deferred task resumption in ``dag.test()`` (#51182) (#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 1b83f7135d4b5167972ec32752642924e8e0a55a)
Co-authored-by: Kaxil Naik <[email protected]>
---
airflow-core/tests/unit/cli/commands/test_dag_command.py | 2 +-
task-sdk/src/airflow/sdk/definitions/dag.py | 7 ++++++-
2 files changed, 7 insertions(+), 2 deletions(-)
diff --git a/airflow-core/tests/unit/cli/commands/test_dag_command.py
b/airflow-core/tests/unit/cli/commands/test_dag_command.py
index fb2daebf418..219785fc7d0 100644
--- a/airflow-core/tests/unit/cli/commands/test_dag_command.py
+++ b/airflow-core/tests/unit/cli/commands/test_dag_command.py
@@ -874,7 +874,7 @@ class TestCliDags:
self.defer(trigger=trigger, method_name="execute")
return
print("RESUMING")
- return self.tfield + 1
+ assert self.tfield + 1 == 3
task_one = one()
task_two = two(task_one)
diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py
b/task-sdk/src/airflow/sdk/definitions/dag.py
index 79616f43ffd..70db3a2035e 100644
--- a/task-sdk/src/airflow/sdk/definitions/dag.py
+++ b/task-sdk/src/airflow/sdk/definitions/dag.py
@@ -1253,6 +1253,8 @@ def _run_task(*, ti, run_triggerer=False):
ti.task = taskrun_result.ti.task
if ti.state == State.DEFERRED and isinstance(msg, DeferTask) and
run_triggerer:
+ from airflow.utils.session import create_session
+
# API Server expects the task instance to be in QUEUED state
before
# resuming from deferral.
ti.set_state(State.QUEUED)
@@ -1264,7 +1266,10 @@ def _run_task(*, ti, run_triggerer=False):
ti.next_kwargs = {"event": event.payload} if event else
msg.next_kwargs
log.info("[DAG TEST] Trigger completed")
- ti.set_state(State.SUCCESS)
+ # Set the state to SCHEDULED so that the task can be resumed.
+ with create_session() as session:
+ ti.state = State.SCHEDULED
+ session.add(ti)
return taskrun_result
except Exception:
