Subscribe also to EVENT_STATE_CHANGED

This commit is contained in:
Erik 2024-06-26 11:48:15 +02:00
parent 2058777db6
commit 9cb8b5be54

View File

@ -27,6 +27,7 @@ from homeassistant.const import (
CONF_METHOD, CONF_METHOD,
CONF_NAME, CONF_NAME,
CONF_UNIQUE_ID, CONF_UNIQUE_ID,
EVENT_STATE_CHANGED,
EVENT_STATE_REPORTED, EVENT_STATE_REPORTED,
STATE_UNAVAILABLE, STATE_UNAVAILABLE,
UnitOfTime, UnitOfTime,
@ -34,6 +35,7 @@ from homeassistant.const import (
from homeassistant.core import ( from homeassistant.core import (
CALLBACK_TYPE, CALLBACK_TYPE,
Event, Event,
EventStateChangedData,
EventStateReportedData, EventStateReportedData,
HomeAssistant, HomeAssistant,
State, State,
@ -435,6 +437,16 @@ class IntegrationSensor(RestoreSensor):
) and state.state != STATE_UNAVAILABLE: ) and state.state != STATE_UNAVAILABLE:
self._derive_and_set_attributes_from_state(state) self._derive_and_set_attributes_from_state(state)
self.async_on_remove(
self.hass.bus.async_listen(
EVENT_STATE_CHANGED,
handle_state_change,
event_filter=callback(
lambda event_data: event_data["entity_id"] == self._sensor_source_id
),
run_immediately=True,
)
)
self.async_on_remove( self.async_on_remove(
self.hass.bus.async_listen( self.hass.bus.async_listen(
EVENT_STATE_REPORTED, EVENT_STATE_REPORTED,
@ -448,7 +460,7 @@ class IntegrationSensor(RestoreSensor):
@callback @callback
def _integrate_on_state_change_and_max_sub_interval( def _integrate_on_state_change_and_max_sub_interval(
self, event: Event[EventStateReportedData] self, event: Event[EventStateChangedData] | Event[EventStateReportedData]
) -> None: ) -> None:
"""Integrate based on state change and time. """Integrate based on state change and time.
@ -456,8 +468,17 @@ class IntegrationSensor(RestoreSensor):
reschedules time based integration. reschedules time based integration.
""" """
self._cancel_max_sub_interval_exceeded_callback() self._cancel_max_sub_interval_exceeded_callback()
old_last_reported = event.data.get("old_last_reported") if event.event_type == EVENT_STATE_CHANGED:
old_state = event.data.get("old_state") if TYPE_CHECKING:
assert type(event.data) is EventStateChangedData
old_last_reported = None
old_state = event.data["old_state"]
else: # EVENT_STATE_REPORTED
if TYPE_CHECKING:
assert type(event.data) is EventStateReportedData
old_last_reported = event.data["old_last_reported"]
old_state = None
new_state = event.data["new_state"] new_state = event.data["new_state"]
try: try:
self._integrate_on_state_change(old_last_reported, old_state, new_state) self._integrate_on_state_change(old_last_reported, old_state, new_state)
@ -470,11 +491,19 @@ class IntegrationSensor(RestoreSensor):
@callback @callback
def _integrate_on_state_change_callback( def _integrate_on_state_change_callback(
self, event: Event[EventStateReportedData] self, event: Event[EventStateChangedData] | Event[EventStateReportedData]
) -> None: ) -> None:
"""Handle the sensor state changes.""" """Handle the sensor state changes."""
old_last_reported = event.data.get("old_last_reported") if event.event_type == EVENT_STATE_CHANGED:
old_state = event.data.get("old_state") if TYPE_CHECKING:
assert type(event.data) is EventStateChangedData
old_last_reported = None
old_state = event.data["old_state"]
else: # EVENT_STATE_REPORTED
if TYPE_CHECKING:
assert type(event.data) is EventStateReportedData
old_last_reported = event.data["old_last_reported"]
old_state = None
new_state = event.data["new_state"] new_state = event.data["new_state"]
return self._integrate_on_state_change(old_last_reported, old_state, new_state) return self._integrate_on_state_change(old_last_reported, old_state, new_state)