diff --git a/esphome/coroutine.py b/esphome/coroutine.py index 3ce94cc9791..8a825362270 100644 --- a/esphome/coroutine.py +++ b/esphome/coroutine.py @@ -45,6 +45,7 @@ the last `yield` expression defines what is returned. from __future__ import annotations from collections.abc import Awaitable, Callable, Generator, Iterator +import contextvars import enum import functools import heapq @@ -277,14 +278,22 @@ class _Task: id_number: int, iterator: Iterator[None], original_function: Any, + context: contextvars.Context, ): self.priority = priority self.id_number = id_number self.iterator = iterator self.original_function = original_function + self.context = context def with_priority(self, priority: float) -> _Task: - return _Task(priority, self.id_number, self.iterator, self.original_function) + return _Task( + priority, + self.id_number, + self.iterator, + self.original_function, + self.context, + ) @property def _cmp_tuple(self) -> tuple[float, int]: @@ -321,7 +330,10 @@ class FakeEventLoop: coro = coroutine(func) gen = coro(*args, **kwargs) prio = getattr(coro, "priority", 0.0) - task = _Task(prio, self._task_counter, gen, func) + # Each task gets its own copy of the current context, isolating any + # contextvars it sets from other tasks the scheduler interleaves it with + # (mirrors what asyncio.Task does internally). + task = _Task(prio, self._task_counter, gen, func, contextvars.copy_context()) self._task_counter += 1 heapq.heappush(self._pending_tasks, task) @@ -352,7 +364,7 @@ class FakeEventLoop: ) try: - next(task.iterator) + task.context.run(next, task.iterator) # Decrease priority over time, so that if this task is blocked # due to a dependency others will clear the dependency # This could be improved with a less naive approach diff --git a/tests/unit_tests/test_coroutine.py b/tests/unit_tests/test_coroutine.py index e12c273294b..0a8fb59cb81 100644 --- a/tests/unit_tests/test_coroutine.py +++ b/tests/unit_tests/test_coroutine.py @@ -1,5 +1,7 @@ """Tests for the coroutine module.""" +import contextvars + import pytest from esphome.coroutine import CoroPriority, FakeEventLoop, coroutine_with_priority @@ -217,3 +219,46 @@ def test_custom_priority_between_enum_values() -> None: # Check execution order assert execution_order == ["core", "custom", "diagnostics"] + + +def test_context_isolated_between_interleaved_tasks() -> None: + """Test that a contextvar set in one task does not leak into another task that the scheduler interleaves with it.""" + my_var: contextvars.ContextVar[str] = contextvars.ContextVar("my_var") + seen: dict[str, str] = {} + + def task_a(): + my_var.set("a") + yield # suspend so task_b can run before task_a resumes + seen["a"] = my_var.get() + + def task_b(): + my_var.set("b") + yield + seen["b"] = my_var.get() + + loop = FakeEventLoop() + loop.add_job(task_a) + loop.add_job(task_b) + loop.flush_tasks() + + assert seen == {"a": "a", "b": "b"} + + +def test_context_inherits_ambient_value_at_schedule_time() -> None: + """Test that a job sees whatever contextvar value was set before it was scheduled.""" + my_var: contextvars.ContextVar[str] = contextvars.ContextVar("my_var") + token = my_var.set("ambient") + seen: dict[str, str] = {} + + def task(): + seen["value"] = my_var.get() + yield + + try: + loop = FakeEventLoop() + loop.add_job(task) + loop.flush_tasks() + finally: + my_var.reset(token) + + assert seen == {"value": "ambient"}