Fix units not being pulled from source sensor (#63639)

homeassistant/components/integration/sensor.py

@@ -150,17 +150,27 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
             old_state = event.data.get("old_state")
             new_state = event.data.get("new_state")
 
+            # We may want to update our state before an early return,
+            # based on the source sensor's unit_of_measurement
+            # or device_class.
+            update_state = False
+
             if self._unit_of_measurement is None:
                 unit = new_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
-                self._unit_of_measurement = self._unit_template.format(
+                if unit is not None:
-                    "" if unit is None else unit
+                    self._unit_of_measurement = self._unit_template.format(unit)
-                )
+                    update_state = True
+
             if (
                 self.device_class is None
                 and new_state.attributes.get(ATTR_DEVICE_CLASS)
                 == SensorDeviceClass.POWER
             ):
                 self._attr_device_class = SensorDeviceClass.ENERGY
+                update_state = True
+
+            if update_state:
+                self.async_write_ha_state()
 
             if (
                 old_state is None

tests/components/integration/test_sensor.py

@@ -3,7 +3,13 @@ from datetime import timedelta
 from unittest.mock import patch
 
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
-from homeassistant.const import ENERGY_KILO_WATT_HOUR, POWER_WATT, TIME_SECONDS
+from homeassistant.const import (
+    ENERGY_KILO_WATT_HOUR,
+    ENERGY_WATT_HOUR,
+    POWER_WATT,
+    STATE_UNKNOWN,
+    TIME_SECONDS,
+)
 from homeassistant.core import HomeAssistant, State
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
@@ -287,3 +293,116 @@ async def test_suffix(hass):
 
     # Testing a network speed sensor at 1000 bytes/s over 10s  = 10kbytes
     assert round(float(state.state)) == 10
+
+
+async def test_units(hass):
+    """Test integration sensor units using a power source."""
+    config = {
+        "sensor": {
+            "platform": "integration",
+            "name": "integration",
+            "source": "sensor.power",
+        }
+    }
+
+    assert await async_setup_component(hass, "sensor", config)
+
+    entity_id = config["sensor"]["source"]
+    # This replicates the current sequence when HA starts up in a real runtime
+    # by updating the base sensor state before the base sensor's units
+    # or state have been correctly populated.  Those interim updates
+    # include states of None and Unknown
+    hass.states.async_set(entity_id, 100, {"unit_of_measurement": None})
+    await hass.async_block_till_done()
+    hass.states.async_set(entity_id, 200, {"unit_of_measurement": None})
+    await hass.async_block_till_done()
+    hass.states.async_set(entity_id, 300, {"unit_of_measurement": POWER_WATT})
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.integration")
+    assert state is not None
+
+    # Testing the sensor ignored the source sensor's units until
+    # they became valid
+    assert state.attributes.get("unit_of_measurement") == ENERGY_WATT_HOUR
+
+
+async def test_device_class(hass):
+    """Test integration sensor units using a power source."""
+    config = {
+        "sensor": {
+            "platform": "integration",
+            "name": "integration",
+            "source": "sensor.power",
+        }
+    }
+
+    assert await async_setup_component(hass, "sensor", config)
+
+    entity_id = config["sensor"]["source"]
+    # This replicates the current sequence when HA starts up in a real runtime
+    # by updating the base sensor state before the base sensor's units
+    # or state have been correctly populated.  Those interim updates
+    # include states of None and Unknown
+    hass.states.async_set(entity_id, STATE_UNKNOWN, {})
+    await hass.async_block_till_done()
+    hass.states.async_set(entity_id, 100, {"device_class": None})
+    await hass.async_block_till_done()
+    hass.states.async_set(entity_id, 200, {"device_class": None})
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.integration")
+    assert "device_class" not in state.attributes
+
+    hass.states.async_set(
+        entity_id, 300, {"device_class": SensorDeviceClass.POWER}, force_update=True
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.integration")
+    assert state is not None
+    # Testing the sensor ignored the source sensor's device class until
+    # it became valid
+    assert state.attributes.get("device_class") == SensorDeviceClass.ENERGY
+
+
+async def test_calc_errors(hass):
+    """Test integration sensor units using a power source."""
+    config = {
+        "sensor": {
+            "platform": "integration",
+            "name": "integration",
+            "source": "sensor.power",
+        }
+    }
+
+    assert await async_setup_component(hass, "sensor", config)
+
+    entity_id = config["sensor"]["source"]
+
+    hass.states.async_set(entity_id, None, {})
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.integration")
+    # With the source sensor in a None state, the Reimann sensor should be
+    # unknown
+    assert state is not None
+    assert state.state == STATE_UNKNOWN
+
+    # Moving from an unknown state to a value is a calc error and should
+    # not change the value of the Reimann sensor.
+    hass.states.async_set(entity_id, 0, {"device_class": None})
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.integration")
+    assert state is not None
+    assert state.state == STATE_UNKNOWN
+
+    # With the source sensor updated successfully, the Reimann sensor
+    # should have a zero (known) value.
+    hass.states.async_set(entity_id, 1, {"device_class": None})
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.integration")
+    assert state is not None
+    assert round(float(state.state)) == 0