From a5a631148ec81c1ca1fb0ec9acb4d78670da908b Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Wed, 26 Jun 2024 22:04:27 -0500
Subject: [PATCH] Add async_track_state_reported_event to fix integration
 performance regression (#120622)

split from https://github.com/home-assistant/core/pull/120621
---
 homeassistant/helpers/event.py | 37 ++++++++++++++++++++++++++++------
 tests/helpers/test_event.py    | 32 ++++++++++++++++++++++++++++-
 2 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py
index 4150d871b6b..ebd51948e3b 100644
--- a/homeassistant/helpers/event.py
+++ b/homeassistant/helpers/event.py
@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Concatenate, Generic, TypeVar
 from homeassistant.const import (
     EVENT_CORE_CONFIG_UPDATE,
     EVENT_STATE_CHANGED,
+    EVENT_STATE_REPORTED,
     MATCH_ALL,
     SUN_EVENT_SUNRISE,
     SUN_EVENT_SUNSET,
@@ -26,6 +27,7 @@ from homeassistant.core import (
     Event,
     # Explicit reexport of 'EventStateChangedData' for backwards compatibility
     EventStateChangedData as EventStateChangedData,  # noqa: PLC0414
+    EventStateReportedData,
     HassJob,
     HassJobType,
     HomeAssistant,
@@ -57,6 +59,9 @@ from .typing import TemplateVarsType
 _TRACK_STATE_CHANGE_DATA: HassKey[_KeyedEventData[EventStateChangedData]] = HassKey(
     "track_state_change_data"
 )
+_TRACK_STATE_REPORTED_DATA: HassKey[_KeyedEventData[EventStateReportedData]] = HassKey(
+    "track_state_reported_data"
+)
 _TRACK_STATE_ADDED_DOMAIN_DATA: HassKey[_KeyedEventData[EventStateChangedData]] = (
     HassKey("track_state_added_domain_data")
 )
@@ -324,8 +329,8 @@ def async_track_state_change_event(
 @callback
 def _async_dispatch_entity_id_event(
     hass: HomeAssistant,
-    callbacks: dict[str, list[HassJob[[Event[EventStateChangedData]], Any]]],
-    event: Event[EventStateChangedData],
+    callbacks: dict[str, list[HassJob[[Event[_TypedDictT]], Any]]],
+    event: Event[_TypedDictT],
 ) -> None:
     """Dispatch to listeners."""
     if not (callbacks_list := callbacks.get(event.data["entity_id"])):
@@ -342,10 +347,10 @@ def _async_dispatch_entity_id_event(
 
 
 @callback
-def _async_state_change_filter(
+def _async_state_filter(
     hass: HomeAssistant,
-    callbacks: dict[str, list[HassJob[[Event[EventStateChangedData]], Any]]],
-    event_data: EventStateChangedData,
+    callbacks: dict[str, list[HassJob[[Event[_TypedDictT]], Any]]],
+    event_data: _TypedDictT,
 ) -> bool:
     """Filter state changes by entity_id."""
     return event_data["entity_id"] in callbacks
@@ -355,7 +360,7 @@ _KEYED_TRACK_STATE_CHANGE = _KeyedEventTracker(
     key=_TRACK_STATE_CHANGE_DATA,
     event_type=EVENT_STATE_CHANGED,
     dispatcher_callable=_async_dispatch_entity_id_event,
-    filter_callable=_async_state_change_filter,
+    filter_callable=_async_state_filter,
 )
 
 
@@ -372,6 +377,26 @@ def _async_track_state_change_event(
     )
 
 
+_KEYED_TRACK_STATE_REPORTED = _KeyedEventTracker(
+    key=_TRACK_STATE_REPORTED_DATA,
+    event_type=EVENT_STATE_REPORTED,
+    dispatcher_callable=_async_dispatch_entity_id_event,
+    filter_callable=_async_state_filter,
+)
+
+
+def async_track_state_reported_event(
+    hass: HomeAssistant,
+    entity_ids: str | Iterable[str],
+    action: Callable[[Event[EventStateReportedData]], Any],
+    job_type: HassJobType | None = None,
+) -> CALLBACK_TYPE:
+    """Track EVENT_STATE_REPORTED by entity_id without lowercasing."""
+    return _async_track_event(
+        _KEYED_TRACK_STATE_REPORTED, hass, entity_ids, action, job_type
+    )
+
+
 @callback
 def _remove_empty_listener() -> None:
     """Remove a listener that does nothing."""
diff --git a/tests/helpers/test_event.py b/tests/helpers/test_event.py
index edce36218e8..4f983120e36 100644
--- a/tests/helpers/test_event.py
+++ b/tests/helpers/test_event.py
@@ -15,7 +15,13 @@ import pytest
 
 from homeassistant.const import MATCH_ALL
 import homeassistant.core as ha
-from homeassistant.core import Event, EventStateChangedData, HomeAssistant, callback
+from homeassistant.core import (
+    Event,
+    EventStateChangedData,
+    EventStateReportedData,
+    HomeAssistant,
+    callback,
+)
 from homeassistant.exceptions import TemplateError
 from homeassistant.helpers.device_registry import EVENT_DEVICE_REGISTRY_UPDATED
 from homeassistant.helpers.entity_registry import EVENT_ENTITY_REGISTRY_UPDATED
@@ -34,6 +40,7 @@ from homeassistant.helpers.event import (
     async_track_state_change_event,
     async_track_state_change_filtered,
     async_track_state_removed_domain,
+    async_track_state_reported_event,
     async_track_sunrise,
     async_track_sunset,
     async_track_template,
@@ -4907,3 +4914,26 @@ async def test_track_point_in_time_repr(
     assert "Exception in callback _TrackPointUTCTime" in caplog.text
     assert "._raise_exception" in caplog.text
     await hass.async_block_till_done(wait_background_tasks=True)
+
+
+async def test_async_track_state_reported_event(hass: HomeAssistant) -> None:
+    """Test async_track_state_reported_event."""
+    tracker_called: list[ha.State] = []
+
+    @ha.callback
+    def single_run_callback(event: Event[EventStateReportedData]) -> None:
+        new_state = event.data["new_state"]
+        tracker_called.append(new_state)
+
+    unsub = async_track_state_reported_event(
+        hass, ["light.bowl", "light.top"], single_run_callback
+    )
+    hass.states.async_set("light.bowl", "on")
+    hass.states.async_set("light.top", "on")
+    await hass.async_block_till_done()
+    assert len(tracker_called) == 0
+    hass.states.async_set("light.bowl", "on")
+    hass.states.async_set("light.top", "on")
+    await hass.async_block_till_done()
+    assert len(tracker_called) == 2
+    unsub()