Small performance improvements to energy sensors (#114149)

homeassistant/components/energy/sensor.py

@@ -149,7 +149,7 @@ class SensorManager:
         async def finish() -> None:
             if to_add:
                 self.async_add_entities(to_add)
-                await asyncio.gather(*(ent.add_finished.wait() for ent in to_add))
+                await asyncio.wait(ent.add_finished for ent in to_add)
 
             for key, entity in to_remove.items():
                 self.current_entities.pop(key)
@@ -219,6 +219,12 @@ class SensorManager:
         to_add.append(self.current_entities[key])
 
 
+def _set_result_unless_done(future: asyncio.Future[None]) -> None:
+    """Set the result of a future unless it is done."""
+    if not future.done():
+        future.set_result(None)
+
+
 class EnergyCostSensor(SensorEntity):
     """Calculate costs incurred by consuming energy.
 
@@ -248,7 +254,9 @@ class EnergyCostSensor(SensorEntity):
         self._last_energy_sensor_state: State | None = None
         # add_finished is set when either of async_added_to_hass or add_to_platform_abort
         # is called
-        self.add_finished = asyncio.Event()
+        self.add_finished: asyncio.Future[
+            None
+        ] = asyncio.get_running_loop().create_future()
 
     def _reset(self, energy_state: State) -> None:
         """Reset the cost sensor."""
@@ -419,25 +427,25 @@ class EnergyCostSensor(SensorEntity):
             self._config[self._adapter.stat_energy_key]
         ] = self.entity_id
 
-        @callback
-        def async_state_changed_listener(*_: Any) -> None:
-            """Handle child updates."""
-            self._update_cost()
-            self.async_write_ha_state()
-
         self.async_on_remove(
             async_track_state_change_event(
                 self.hass,
                 cast(str, self._config[self._adapter.stat_energy_key]),
-                async_state_changed_listener,
+                self._async_state_changed_listener,
             )
         )
-        self.add_finished.set()
+        _set_result_unless_done(self.add_finished)
+
+    @callback
+    def _async_state_changed_listener(self, *_: Any) -> None:
+        """Handle child updates."""
+        self._update_cost()
+        self.async_write_ha_state()
 
     @callback
     def add_to_platform_abort(self) -> None:
         """Abort adding an entity to a platform."""
-        self.add_finished.set()
+        _set_result_unless_done(self.add_finished)
         super().add_to_platform_abort()
 
     async def async_will_remove_from_hass(self) -> None: