(airflow) branch main updated: Fix dag.test() dropping defer-time kwargs when trigger yields an event (#69942)

[vi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 4be96fbd96a Fix dag.test() dropping defer-time kwargs when trigger yields an event (#69942)
4be96fbd96a is described below

commit 4be96fbd96a0b1d5717b705bb970c7fc67f2b0bd
Author: Sean Ghaeli <58...@users.noreply.github.com>
AuthorDate: Mon Jul 20 13:51:00 2026 -0700

    Fix dag.test() dropping defer-time kwargs when trigger yields an event (#69942)
    
    When DAG.test() runs a deferred task's trigger inline and the trigger
    yields an event, the kwargs passed at defer() time were discarded and
    replaced with only {"event": ...}. The resume method then ran with its
    default argument values.
    
    This broke example_neptune_analytics on the deferrable path:
    NeptuneCreatePrivateGraphEndpointOperator defers with
    kwargs={"vpc_id": self.vpc_id} and resumes into
    execute_complete(..., vpc_id=""). The lost vpc_id fell back to "" and
    was pushed to XCom; the downstream delete_endpoint task built a malformed
    request URI (/graphs/<id>/endpoints/) and the service rejected it with
    AccessDeniedException: Unable to determine service/operation name to be
    authorized.
    
    Merge the event into the deferral's next_kwargs instead of replacing
    them, mirroring the production resume path in airflow.models.trigger.
    Add a regression test.
---
 airflow-core/tests/unit/models/test_dag.py  | 27 +++++++++++++++++++++++++++
 task-sdk/src/airflow/sdk/definitions/dag.py | 11 ++++++++++-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py
index 25d02d86686..ead765b3171 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -1958,6 +1958,33 @@ class TestDag:
         assert parked_states_seen == [TaskInstanceState.AWAITING_INPUT]
         assert resume_calls == [(["Approve"], {})]
 
+    @pytest.mark.execution_timeout(60)
+    def test_dag_test_preserves_defer_kwargs_through_inline_trigger(self, testing_dag_bundle):
+        """dag.test() must keep defer()-time kwargs when the inline trigger yields an event."""
+        from airflow.providers.standard.triggers.temporal import TimeDeltaTrigger
+
+        seen_kwargs: list = []
+
+        class DeferKwargOperator(BaseOperator):
+            def execute(self, context):
+                self.defer(
+                    trigger=TimeDeltaTrigger(datetime.timedelta(seconds=0)),
+                    method_name="resume",
+                    kwargs={"vpc_id": "vpc-abc123"},
+                )
+
+            def resume(self, context, event=None, vpc_id=""):
+                seen_kwargs.append(vpc_id)
+
+        dag = DAG(dag_id="test_dag_test_defer_kwargs", schedule=None, start_date=DEFAULT_DATE)
+        with dag:
+            DeferKwargOperator(task_id="defer_task")
+        sync_dag_to_db(dag)
+
+        dag.test()
+
+        assert seen_kwargs == ["vpc-abc123"]  # was [""] before the fix
+
     def test_dag_connection_file(self, tmp_path, testing_dag_bundle):
         test_connections_string = """
 ---
diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py
index 5ba0254b208..cadfe697830 100644
--- a/task-sdk/src/airflow/sdk/definitions/dag.py
+++ b/task-sdk/src/airflow/sdk/definitions/dag.py
@@ -1544,7 +1544,16 @@ def _run_task(
                 trigger = import_string(msg.classpath)(**kwargs)
                 event = _run_inline_trigger(trigger, task_sdk_ti)
                 ti.next_method = msg.next_method
-                ti.next_kwargs = {"event": serialize(event.payload)} if event else msg.next_kwargs
+                # Add the trigger event to the kwargs the task passed to defer(), instead of
+                # discarding them, so the resume method still receives those kwargs.
+                if event:
+                    next_kwargs = deserialize(msg.next_kwargs) if msg.next_kwargs else {}
+                    if TYPE_CHECKING:
+                        assert isinstance(next_kwargs, dict)
+                    next_kwargs["event"] = event.payload
+                    ti.next_kwargs = serialize(next_kwargs)
+                else:
+                    ti.next_kwargs = msg.next_kwargs
                 log.info("[DAG TEST] Trigger completed")
 
                 # Set the state to SCHEDULED so that the task can be resumed.