diff --git a/homeassistant/components/energy/sensor.py b/homeassistant/components/energy/sensor.py index 638b560a954..cca245f7e8d 100644 --- a/homeassistant/components/energy/sensor.py +++ b/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: