From b15ea588510601789af555b6beb3abf6639f1a25 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 12 Aug 2024 02:15:33 -0500
Subject: [PATCH] Relocate code to get scheduled TimerHandles (#123546)

---
 homeassistant/core.py        |  4 ++--
 homeassistant/util/async_.py | 16 +++++++++++++++-
 tests/common.py              |  3 ++-
 tests/conftest.py            |  4 ++--
 tests/util/test_async.py     | 14 ++++++++++++++
 5 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/homeassistant/core.py b/homeassistant/core.py
index 5d223b9f19f..1050d25ee71 100644
--- a/homeassistant/core.py
+++ b/homeassistant/core.py
@@ -101,6 +101,7 @@ from .util import dt as dt_util, location
 from .util.async_ import (
     cancelling,
     create_eager_task,
+    get_scheduled_timer_handles,
     run_callback_threadsafe,
     shutdown_run_callback_threadsafe,
 )
@@ -1227,8 +1228,7 @@ class HomeAssistant:
 
     def _cancel_cancellable_timers(self) -> None:
         """Cancel timer handles marked as cancellable."""
-        handles: Iterable[asyncio.TimerHandle] = self.loop._scheduled  # type: ignore[attr-defined] # noqa: SLF001
-        for handle in handles:
+        for handle in get_scheduled_timer_handles(self.loop):
             if (
                 not handle.cancelled()
                 and (args := handle._args)  # noqa: SLF001
diff --git a/homeassistant/util/async_.py b/homeassistant/util/async_.py
index f2dc1291324..dcb788f0685 100644
--- a/homeassistant/util/async_.py
+++ b/homeassistant/util/async_.py
@@ -2,7 +2,15 @@
 
 from __future__ import annotations
 
-from asyncio import AbstractEventLoop, Future, Semaphore, Task, gather, get_running_loop
+from asyncio import (
+    AbstractEventLoop,
+    Future,
+    Semaphore,
+    Task,
+    TimerHandle,
+    gather,
+    get_running_loop,
+)
 from collections.abc import Awaitable, Callable, Coroutine
 import concurrent.futures
 import logging
@@ -124,3 +132,9 @@ def shutdown_run_callback_threadsafe(loop: AbstractEventLoop) -> None:
     python is going to exit.
     """
     setattr(loop, _SHUTDOWN_RUN_CALLBACK_THREADSAFE, True)
+
+
+def get_scheduled_timer_handles(loop: AbstractEventLoop) -> list[TimerHandle]:
+    """Return a list of scheduled TimerHandles."""
+    handles: list[TimerHandle] = loop._scheduled  # type: ignore[attr-defined] # noqa: SLF001
+    return handles
diff --git a/tests/common.py b/tests/common.py
index 64e11ee7b51..d36df509142 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -93,6 +93,7 @@ from homeassistant.helpers.json import JSONEncoder, _orjson_default_encoder, jso
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 from homeassistant.util.async_ import (
     _SHUTDOWN_RUN_CALLBACK_THREADSAFE,
+    get_scheduled_timer_handles,
     run_callback_threadsafe,
 )
 import homeassistant.util.dt as dt_util
@@ -531,7 +532,7 @@ def _async_fire_time_changed(
     hass: HomeAssistant, utc_datetime: datetime | None, fire_all: bool
 ) -> None:
     timestamp = dt_util.utc_to_timestamp(utc_datetime)
-    for task in list(hass.loop._scheduled):
+    for task in list(get_scheduled_timer_handles(hass.loop)):
         if not isinstance(task, asyncio.TimerHandle):
             continue
         if task.cancelled():
diff --git a/tests/conftest.py b/tests/conftest.py
index 0667edf4be2..ea0453e7450 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -84,7 +84,7 @@ from homeassistant.helpers.translation import _TranslationsCacheData
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.setup import async_setup_component
 from homeassistant.util import dt as dt_util, location
-from homeassistant.util.async_ import create_eager_task
+from homeassistant.util.async_ import create_eager_task, get_scheduled_timer_handles
 from homeassistant.util.json import json_loads
 
 from .ignore_uncaught_exceptions import IGNORE_UNCAUGHT_EXCEPTIONS
@@ -372,7 +372,7 @@ def verify_cleanup(
     if tasks:
         event_loop.run_until_complete(asyncio.wait(tasks))
 
-    for handle in event_loop._scheduled:  # type: ignore[attr-defined]
+    for handle in get_scheduled_timer_handles(event_loop):
         if not handle.cancelled():
             with long_repr_strings():
                 if expected_lingering_timers:
diff --git a/tests/util/test_async.py b/tests/util/test_async.py
index 373768788b7..17349cf6ff9 100644
--- a/tests/util/test_async.py
+++ b/tests/util/test_async.py
@@ -199,3 +199,17 @@ async def test_create_eager_task_from_thread_in_integration(
         "from a thread at homeassistant/components/hue/light.py, line 23: "
         "self.light.is_on"
     ) in caplog.text
+
+
+async def test_get_scheduled_timer_handles(hass: HomeAssistant) -> None:
+    """Test get_scheduled_timer_handles returns all scheduled timer handles."""
+    loop = hass.loop
+    timer_handle = loop.call_later(10, lambda: None)
+    timer_handle2 = loop.call_later(5, lambda: None)
+    timer_handle3 = loop.call_later(15, lambda: None)
+
+    handles = hasync.get_scheduled_timer_handles(loop)
+    assert set(handles).issuperset({timer_handle, timer_handle2, timer_handle3})
+    timer_handle.cancel()
+    timer_handle2.cancel()
+    timer_handle3.cancel()