dstandish commented on code in PR #69146:
URL: https://github.com/apache/airflow/pull/69146#discussion_r3509482599
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -5597,6 +5676,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"
+ task.on_kill.assert_not_called()
+
+ def test_applies_execution_timeout(self):
+ """When a timeout is set and the callable overruns, AirflowTaskTimeout
is raised and on_kill is called."""
+ task = self._make_task(execution_timeout=timedelta(milliseconds=10))
+
+ def execute(context):
+ time.sleep(2)
Review Comment:
good call
--
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]
