Refactor event handlers

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(
+    def _integrate_on_state_change_with_max_sub_interval(
-        self, event: Event[EventStateChangedData] | Event[EventStateReportedData]
+        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."""
+        """Handle sensor state change."""
-        if event.event_type == EVENT_STATE_CHANGED:
+        return self._integrate_on_state_change(
-            if TYPE_CHECKING:
+            None, event.data["old_state"], event.data["new_state"]
-                assert type(event.data) is EventStateChangedData
+        )
-            old_last_reported = None
+
-            old_state = event.data["old_state"]
+    @callback
-        else:  # EVENT_STATE_REPORTED
+    def _integrate_on_state_report_callback(
-            if TYPE_CHECKING:
+        self, event: Event[EventStateReportedData]
-                assert type(event.data) is EventStateReportedData
+    ) -> None:
-            old_last_reported = event.data["old_last_reported"]
+        """Handle sensor state report."""
-            old_state = None
+        return self._integrate_on_state_change(
-        new_state = event.data["new_state"]
+            event.data["old_last_reported"], None, event.data["new_state"]
-        return self._integrate_on_state_change(old_last_reported, old_state, 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()
         )