Define unit_of_measurement of all utility_meter sensors on HA start (#56112)

* define unit_of_measurement on hass start

* delay utility_meter state

* check state

* store siblings

* don't check unit_of_measurement

homeassistant/components/utility_meter/__init__.py

@@ -24,6 +24,7 @@ from .const import (
     CONF_TARIFF,
     CONF_TARIFF_ENTITY,
     CONF_TARIFFS,
+    DATA_TARIFF_SENSORS,
     DATA_UTILITY,
     DOMAIN,
     METER_TYPES,
@@ -98,6 +99,7 @@ async def async_setup(hass, config):
         _LOGGER.debug("Setup %s.%s", DOMAIN, meter)
 
         hass.data[DATA_UTILITY][meter] = conf
+        hass.data[DATA_UTILITY][meter][DATA_TARIFF_SENSORS] = []
 
         if not conf[CONF_TARIFFS]:
             # only one entity is required

homeassistant/components/utility_meter/const.py

@@ -22,6 +22,7 @@ METER_TYPES = [
 ]
 
 DATA_UTILITY = "utility_meter_data"
+DATA_TARIFF_SENSORS = "utility_meter_sensors"
 
 CONF_METER = "meter"
 CONF_SOURCE_SENSOR = "source"

homeassistant/components/utility_meter/sensor.py

@@ -46,6 +46,7 @@ from .const import (
     CONF_TARIFF,
     CONF_TARIFF_ENTITY,
     DAILY,
+    DATA_TARIFF_SENSORS,
     DATA_UTILITY,
     HOURLY,
     MONTHLY,
@@ -96,19 +97,20 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
             CONF_TARIFF_ENTITY
         )
         conf_cron_pattern = hass.data[DATA_UTILITY][meter].get(CONF_CRON_PATTERN)
-
+        meter_sensor = UtilityMeterSensor(
-        meters.append(
+            meter,
-            UtilityMeterSensor(
+            conf_meter_source,
-                conf_meter_source,
+            conf.get(CONF_NAME),
-                conf.get(CONF_NAME),
+            conf_meter_type,
-                conf_meter_type,
+            conf_meter_offset,
-                conf_meter_offset,
+            conf_meter_net_consumption,
-                conf_meter_net_consumption,
+            conf.get(CONF_TARIFF),
-                conf.get(CONF_TARIFF),
+            conf_meter_tariff_entity,
-                conf_meter_tariff_entity,
+            conf_cron_pattern,
-                conf_cron_pattern,
-            )
         )
+        meters.append(meter_sensor)
+
+        hass.data[DATA_UTILITY][meter][DATA_TARIFF_SENSORS].append(meter_sensor)
 
     async_add_entities(meters)
 
@@ -126,6 +128,7 @@ class UtilityMeterSensor(RestoreEntity, SensorEntity):
 
     def __init__(
         self,
+        parent_meter,
         source_entity,
         name,
         meter_type,
@@ -136,8 +139,9 @@ class UtilityMeterSensor(RestoreEntity, SensorEntity):
         cron_pattern=None,
     ):
         """Initialize the Utility Meter sensor."""
+        self._parent_meter = parent_meter
         self._sensor_source_id = source_entity
-        self._state = 0
+        self._state = None
         self._last_period = 0
         self._last_reset = dt_util.utcnow()
         self._collecting = None
@@ -153,11 +157,26 @@ class UtilityMeterSensor(RestoreEntity, SensorEntity):
         self._tariff = tariff
         self._tariff_entity = tariff_entity
 
+    def start(self, unit):
+        """Initialize unit and state upon source initial update."""
+        self._unit_of_measurement = unit
+        self._state = 0
+        self.async_write_ha_state()
+
     @callback
     def async_reading(self, event):
         """Handle the sensor state changes."""
         old_state = event.data.get("old_state")
         new_state = event.data.get("new_state")
+
+        if self._state is None and new_state.state:
+            # First state update initializes the utility_meter sensors
+            source_state = self.hass.states.get(self._sensor_source_id)
+            for sensor in self.hass.data[DATA_UTILITY][self._parent_meter][
+                DATA_TARIFF_SENSORS
+            ]:
+                sensor.start(source_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT))
+
         if (
             old_state is None
             or new_state is None
@@ -333,7 +352,12 @@ class UtilityMeterSensor(RestoreEntity, SensorEntity):
                 self._change_status(tariff_entity_state.state)
                 return
 
-            _LOGGER.debug("<%s> collecting from %s", self.name, self._sensor_source_id)
+            _LOGGER.debug(
+                "<%s> collecting %s from %s",
+                self.name,
+                self._unit_of_measurement,
+                self._sensor_source_id,
+            )
             self._collecting = async_track_state_change_event(
                 self.hass, [self._sensor_source_id], self.async_reading
             )

tests/components/utility_meter/test_sensor.py

@@ -31,6 +31,7 @@ from homeassistant.const import (
     ENERGY_KILO_WATT_HOUR,
     EVENT_HOMEASSISTANT_START,
     STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
 )
 from homeassistant.core import State
 from homeassistant.setup import async_setup_component
@@ -64,6 +65,8 @@ async def test_state(hass):
     await hass.async_block_till_done()
 
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
+    await hass.async_block_till_done()
+
     entity_id = config[DOMAIN]["energy_bill"]["source"]
     hass.states.async_set(
         entity_id, 2, {ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR}
@@ -74,16 +77,19 @@ async def test_state(hass):
     assert state is not None
     assert state.state == "0"
     assert state.attributes.get("status") == COLLECTING
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
 
     state = hass.states.get("sensor.energy_bill_midpeak")
     assert state is not None
     assert state.state == "0"
     assert state.attributes.get("status") == PAUSED
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
 
     state = hass.states.get("sensor.energy_bill_offpeak")
     assert state is not None
     assert state.state == "0"
     assert state.attributes.get("status") == PAUSED
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
 
     now = dt_util.utcnow() + timedelta(seconds=10)
     with patch("homeassistant.util.dt.utcnow", return_value=now):
@@ -187,6 +193,49 @@ async def test_state(hass):
     assert state.state == "0.123"
 
 
+async def test_init(hass):
+    """Test utility sensor state initializtion."""
+    config = {
+        "utility_meter": {
+            "energy_bill": {
+                "source": "sensor.energy",
+                "tariffs": ["onpeak", "midpeak", "offpeak"],
+            }
+        }
+    }
+
+    assert await async_setup_component(hass, DOMAIN, config)
+    await hass.async_block_till_done()
+
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.energy_bill_onpeak")
+    assert state is not None
+    assert state.state == STATE_UNKNOWN
+
+    state = hass.states.get("sensor.energy_bill_offpeak")
+    assert state is not None
+    assert state.state == STATE_UNKNOWN
+
+    entity_id = config[DOMAIN]["energy_bill"]["source"]
+    hass.states.async_set(
+        entity_id, 2, {ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR}
+    )
+
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.energy_bill_onpeak")
+    assert state is not None
+    assert state.state == "0"
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
+
+    state = hass.states.get("sensor.energy_bill_offpeak")
+    assert state is not None
+    assert state.state == "0"
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
+
+
 async def test_device_class(hass):
     """Test utility device_class."""
     config = {
@@ -205,6 +254,8 @@ async def test_device_class(hass):
     await hass.async_block_till_done()
 
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
+    await hass.async_block_till_done()
+
     entity_id_energy = config[DOMAIN]["energy_meter"]["source"]
     hass.states.async_set(
         entity_id_energy, 2, {ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR}
@@ -218,35 +269,13 @@ async def test_device_class(hass):
     state = hass.states.get("sensor.energy_meter")
     assert state is not None
     assert state.state == "0"
-    assert state.attributes.get(ATTR_DEVICE_CLASS) is None
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_TOTAL
-    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) is None
-
-    state = hass.states.get("sensor.gas_meter")
-    assert state is not None
-    assert state.state == "0"
-    assert state.attributes.get(ATTR_DEVICE_CLASS) is None
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_TOTAL_INCREASING
-    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) is None
-
-    hass.states.async_set(
-        entity_id_energy, 3, {ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR}
-    )
-    hass.states.async_set(
-        entity_id_gas, 3, {ATTR_UNIT_OF_MEASUREMENT: "some_archaic_unit"}
-    )
-    await hass.async_block_till_done()
-
-    state = hass.states.get("sensor.energy_meter")
-    assert state is not None
-    assert state.state == "1"
     assert state.attributes.get(ATTR_DEVICE_CLASS) == "energy"
     assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_TOTAL
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
 
     state = hass.states.get("sensor.gas_meter")
     assert state is not None
-    assert state.state == "1"
+    assert state.state == "0"
     assert state.attributes.get(ATTR_DEVICE_CLASS) is None
     assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_TOTAL_INCREASING
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "some_archaic_unit"
@@ -272,6 +301,7 @@ async def test_restore_state(hass):
                 attributes={
                     ATTR_STATUS: PAUSED,
                     ATTR_LAST_RESET: last_reset,
+                    ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR,
                 },
             ),
             State(
@@ -280,6 +310,7 @@ async def test_restore_state(hass):
                 attributes={
                     ATTR_STATUS: COLLECTING,
                     ATTR_LAST_RESET: last_reset,
+                    ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR,
                 },
             ),
         ],
@@ -293,11 +324,13 @@ async def test_restore_state(hass):
     assert state.state == "3"
     assert state.attributes.get("status") == PAUSED
     assert state.attributes.get("last_reset") == last_reset
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
 
     state = hass.states.get("sensor.energy_bill_offpeak")
     assert state.state == "6"
     assert state.attributes.get("status") == COLLECTING
     assert state.attributes.get("last_reset") == last_reset
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
 
     # utility_meter is loaded, now set sensors according to utility_meter:
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START)