From 443147e132fd3714396aaf651a80840a0d2cfd1b Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Fri, 10 Sep 2021 18:07:52 +0200
Subject: [PATCH] Wait for entities when updating energy preferences (#56057)

---
 homeassistant/components/energy/sensor.py | 15 +++++++++++++--
 tests/components/energy/test_validate.py  |  3 +++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/energy/sensor.py b/homeassistant/components/energy/sensor.py
index fadec44c1a2..387a08141a2 100644
--- a/homeassistant/components/energy/sensor.py
+++ b/homeassistant/components/energy/sensor.py
@@ -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."""
diff --git a/tests/components/energy/test_validate.py b/tests/components/energy/test_validate.py
index 8c67f3eabaf..e893c71d1f2 100644
--- a/tests/components/energy/test_validate.py
+++ b/tests/components/energy/test_validate.py
@@ -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": [],