Wait for entities when updating energy preferences (#56057)

This commit is contained in:
Erik Montnemery 2021-09-10 18:07:52 +02:00 committed by GitHub
parent c59540cfc7
commit 443147e132
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 2 deletions

View File

@ -1,6 +1,7 @@
"""Helper sensor for calculating utility costs."""
from __future__ import annotations
import asyncio
import copy
from dataclasses import dataclass
import logging
@ -117,12 +118,13 @@ class SensorManager:
async def _process_manager_data(self) -> None:
"""Process manager data."""
to_add: list[SensorEntity] = []
to_add: list[EnergyCostSensor] = []
to_remove = dict(self.current_entities)
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))
for key, entity in to_remove.items():
self.current_entities.pop(key)
@ -163,7 +165,7 @@ class SensorManager:
self,
adapter: SourceAdapter,
config: dict,
to_add: list[SensorEntity],
to_add: list[EnergyCostSensor],
to_remove: dict[tuple[str, str | None, str], EnergyCostSensor],
) -> None:
"""Process sensor data."""
@ -220,6 +222,9 @@ class EnergyCostSensor(SensorEntity):
self._config = config
self._last_energy_sensor_state: State | None = None
self._cur_value = 0.0
# add_finished is set when either of async_added_to_hass or add_to_platform_abort
# is called
self.add_finished = asyncio.Event()
def _reset(self, energy_state: State) -> None:
"""Reset the cost sensor."""
@ -373,6 +378,12 @@ class EnergyCostSensor(SensorEntity):
async_state_changed_listener,
)
)
self.add_finished.set()
@callback
def add_to_platform_abort(self) -> None:
"""Abort adding an entity to a platform."""
self.add_finished.set()
async def async_will_remove_from_hass(self) -> None:
"""Handle removing from hass."""

View File

@ -341,12 +341,14 @@ async def test_validation_grid_price_not_exist(hass, mock_energy_manager):
"stat_energy_from": "sensor.grid_consumption_1",
"entity_energy_from": "sensor.grid_consumption_1",
"entity_energy_price": "sensor.grid_price_1",
"number_energy_price": None,
}
],
"flow_to": [
{
"stat_energy_to": "sensor.grid_production_1",
"entity_energy_to": "sensor.grid_production_1",
"entity_energy_price": None,
"number_energy_price": 0.10,
}
],
@ -417,6 +419,7 @@ async def test_validation_grid_price_errors(
"stat_energy_from": "sensor.grid_consumption_1",
"entity_energy_from": "sensor.grid_consumption_1",
"entity_energy_price": "sensor.grid_price_1",
"number_energy_price": None,
}
],
"flow_to": [],