dstandish commented on code in PR #67877:
URL: https://github.com/apache/airflow/pull/67877#discussion_r3501429026
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -5597,6 +5600,141 @@ def test_exception_in_context_manager_propagates(self):
raise ValueError("boom")
+class TestRunExecuteCallable:
+ """Tests for ``_run_execute_callable``.
+
+ It runs the task's execute callable inside an isolated contextvars copy
(with
+ the ExecutorSafeguard tracker set), applies the execution timeout when one
is
+ configured, and wraps the call in a ``task.execute`` detail span.
+ """
+
+ @pytest.fixture(autouse=True)
+ def _sampled_carrier_provider(self):
+ """Make new_dagrun_trace_carrier produce a SAMPLED carrier (see
TestDetailSpan)."""
+ provider = TracerProvider()
+ with mock.patch(
+ "airflow._shared.observability.traces.trace.get_tracer_provider",
+ return_value=provider,
+ ):
+ yield
+
+ @staticmethod
+ def _make_task(execution_timeout=None):
+ task = mock.MagicMock(spec=BaseOperator)
+ task.execution_timeout = execution_timeout
+ return task
+
+ def test_runs_in_isolated_context_with_safeguard_tracker_set(self):
+ """The callable runs in an internal context copy that has the
safeguard tracker set and does not leak."""
+ var = contextvars.ContextVar("marker")
+ var.set("outer")
+ task = self._make_task()
+ seen = {}
+
+ def execute(context):
+ var.set("inner")
+ seen["tracker"] = ExecutorSafeguard.tracker.get(None)
+ return context["value"] * 2
+
+ result = _run_execute_callable(context={"value": 21}, execute=execute,
task=task)
+
+ assert result == 42
+ # The safeguard tracker is set to the task inside the copy used to run
execute.
+ assert seen["tracker"] is task
+ # The mutation happened inside the copy, so it does not leak to the
caller's context.
+ assert var.get() == "outer"
Review Comment:
Good call — added `assert ExecutorSafeguard.tracker.get(None) is not task`
after the call. You're right that this pins the actual headline behavior (the
`.set` confined to the copy) rather than relying on the `marker` proxy, and `is
not task` over `is None` since `task` is a fresh mock per call and nothing
resets the tracker between tests. Pushed.
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @dstandish before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
