Restores unit_of_measurement (#54335)

homeassistant/components/integration/sensor.py

@@ -145,6 +145,10 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
                 )
                 self._attr_device_class = state.attributes.get(ATTR_DEVICE_CLASS)
 
+                self._unit_of_measurement = state.attributes.get(
+                    ATTR_UNIT_OF_MEASUREMENT
+                )
+
         @callback
         def calc_integration(event):
             """Handle the sensor state changes."""

tests/components/integration/test_sensor.py

@@ -73,6 +73,7 @@ async def test_restore_state(hass: HomeAssistant) -> None:
                 {
                     "last_reset": "2019-10-06T21:00:00",
                     "device_class": DEVICE_CLASS_ENERGY,
+                    "unit_of_measurement": ENERGY_KILO_WATT_HOUR,
                 },
             ),
         ),