From 423f4c5bcfcd0c2dbb6463372c9200200eb95c89 Mon Sep 17 00:00:00 2001 From: Erik Date: Wed, 26 Jun 2024 13:25:52 +0200 Subject: [PATCH] Refactor event handlers --- .../components/integration/sensor.py | 79 +++++++++++-------- 1 file changed, 45 insertions(+), 34 deletions(-) diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py index dae0bcf3ee3..60cbee5549f 100644 --- a/homeassistant/components/integration/sensor.py +++ b/homeassistant/components/integration/sensor.py @@ -178,7 +178,7 @@ _NAME_TO_INTEGRATION_METHOD: dict[str, type[_IntegrationMethod]] = { class _IntegrationTrigger(Enum): - StateReport = "state_report" + StateEvent = "state_event" TimeElapsed = "time_elapsed" @@ -338,7 +338,7 @@ class IntegrationSensor(RestoreSensor): ) self._max_sub_interval_exceeded_callback: CALLBACK_TYPE = lambda *args: None self._last_integration_time: datetime = datetime.now(tz=UTC) - self._last_integration_trigger = _IntegrationTrigger.StateReport + self._last_integration_trigger = _IntegrationTrigger.StateEvent self._attr_suggested_display_precision = round_digits or 2 def _calculate_unit(self, source_unit: str) -> str: @@ -428,9 +428,11 @@ class IntegrationSensor(RestoreSensor): source_state = self.hass.states.get(self._sensor_source_id) self._schedule_max_sub_interval_exceeded_if_state_is_numeric(source_state) self.async_on_remove(self._cancel_max_sub_interval_exceeded_callback) - handle_state_change = self._integrate_on_state_change_and_max_sub_interval + handle_state_change = self._integrate_on_state_change_with_max_sub_interval + handle_state_report = self._integrate_on_state_report_with_max_sub_interval else: handle_state_change = self._integrate_on_state_change_callback + handle_state_report = self._integrate_on_state_report_callback if ( state := self.hass.states.get(self._source_entity) @@ -450,7 +452,7 @@ class IntegrationSensor(RestoreSensor): self.async_on_remove( self.hass.bus.async_listen( EVENT_STATE_REPORTED, - handle_state_change, + handle_state_report, event_filter=callback( lambda event_data: event_data["entity_id"] == self._sensor_source_id ), @@ -459,8 +461,29 @@ class IntegrationSensor(RestoreSensor): ) @callback - def _integrate_on_state_change_and_max_sub_interval( - self, event: Event[EventStateChangedData] | Event[EventStateReportedData] + def _integrate_on_state_change_with_max_sub_interval( + self, event: Event[EventStateChangedData] + ) -> None: + """Handle sensor state update when sub interval is configured.""" + self._integrate_on_state_update_with_max_sub_interval( + None, event.data["old_state"], event.data["new_state"] + ) + + @callback + def _integrate_on_state_report_with_max_sub_interval( + self, event: Event[EventStateReportedData] + ) -> None: + """Handle sensor state report when sub interval is configured.""" + self._integrate_on_state_update_with_max_sub_interval( + event.data["old_last_reported"], None, event.data["new_state"] + ) + + @callback + def _integrate_on_state_update_with_max_sub_interval( + self, + old_last_reported: datetime | None, + old_state: State | None, + new_state: State | None, ) -> None: """Integrate based on state change and time. @@ -468,21 +491,9 @@ class IntegrationSensor(RestoreSensor): reschedules time based integration. """ self._cancel_max_sub_interval_exceeded_callback() - if event.event_type == EVENT_STATE_CHANGED: - 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"] try: self._integrate_on_state_change(old_last_reported, old_state, new_state) - self._last_integration_trigger = _IntegrationTrigger.StateReport + self._last_integration_trigger = _IntegrationTrigger.StateEvent self._last_integration_time = datetime.now(tz=UTC) finally: # When max_sub_interval exceeds without state change the source is assumed @@ -491,21 +502,21 @@ class IntegrationSensor(RestoreSensor): @callback def _integrate_on_state_change_callback( - self, event: Event[EventStateChangedData] | Event[EventStateReportedData] + self, event: Event[EventStateChangedData] ) -> None: - """Handle the sensor state changes.""" - if event.event_type == EVENT_STATE_CHANGED: - 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"] - return self._integrate_on_state_change(old_last_reported, old_state, new_state) + """Handle sensor state change.""" + return self._integrate_on_state_change( + None, event.data["old_state"], event.data["new_state"] + ) + + @callback + def _integrate_on_state_report_callback( + self, event: Event[EventStateReportedData] + ) -> None: + """Handle sensor state report.""" + return self._integrate_on_state_change( + event.data["old_last_reported"], None, event.data["new_state"] + ) def _integrate_on_state_change( self, @@ -546,7 +557,7 @@ class IntegrationSensor(RestoreSensor): assert old_last_reported is not None elapsed_seconds = Decimal( (new_state.last_reported - old_last_reported).total_seconds() - if self._last_integration_trigger == _IntegrationTrigger.StateReport + if self._last_integration_trigger == _IntegrationTrigger.StateEvent else (new_state.last_reported - self._last_integration_time).total_seconds() )