Set device_class and state_class for utility_meter (#50921)

* Set device_class and state_class for utility_meter

* Update test

* Tweak tests according to review comments

homeassistant/components/utility_meter/sensor.py

@@ -5,10 +5,17 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.components.sensor import ATTR_LAST_RESET, SensorEntity
+from homeassistant.components.sensor import (
+    ATTR_LAST_RESET,
+    STATE_CLASS_MEASUREMENT,
+    SensorEntity,
+)
 from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
     CONF_NAME,
+    DEVICE_CLASS_ENERGY,
+    ENERGY_KILO_WATT_HOUR,
+    ENERGY_WATT_HOUR,
     EVENT_HOMEASSISTANT_START,
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
@@ -53,6 +60,11 @@ ATTR_PERIOD = "meter_period"
 ATTR_LAST_PERIOD = "last_period"
 ATTR_TARIFF = "tariff"
 
+DEVICE_CLASS_MAP = {
+    ENERGY_WATT_HOUR: DEVICE_CLASS_ENERGY,
+    ENERGY_KILO_WATT_HOUR: DEVICE_CLASS_ENERGY,
+}
+
 ICON = "mdi:counter"
 
 PRECISION = 3
@@ -313,6 +325,16 @@ class UtilityMeterSensor(RestoreEntity, SensorEntity):
         """Return the state of the sensor."""
         return self._state
 
+    @property
+    def device_class(self):
+        """Return the device class of the sensor."""
+        return DEVICE_CLASS_MAP.get(self.unit_of_measurement)
+
+    @property
+    def state_class(self):
+        """Return the device class of the sensor."""
+        return STATE_CLASS_MEASUREMENT
+
     @property
     def unit_of_measurement(self):
         """Return the unit the value is expressed in."""

tests/components/utility_meter/test_sensor.py

@@ -3,7 +3,7 @@ from contextlib import contextmanager
 from datetime import timedelta
 from unittest.mock import patch
 
-from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
+from homeassistant.components.sensor import ATTR_STATE_CLASS, STATE_CLASS_MEASUREMENT
 from homeassistant.components.utility_meter.const import (
     ATTR_TARIFF,
     ATTR_VALUE,
@@ -18,6 +18,7 @@ from homeassistant.components.utility_meter.sensor import (
     PAUSED,
 )
 from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
     ATTR_UNIT_OF_MEASUREMENT,
     ENERGY_KILO_WATT_HOUR,
@@ -52,7 +53,6 @@ async def test_state(hass):
     }
 
     assert await async_setup_component(hass, DOMAIN, config)
-    assert await async_setup_component(hass, SENSOR_DOMAIN, config)
     await hass.async_block_till_done()
 
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
@@ -159,6 +159,70 @@ async def test_state(hass):
     assert state.state == "0.123"
 
 
+async def test_device_class(hass):
+    """Test utility device_class."""
+    config = {
+        "utility_meter": {
+            "energy_meter": {
+                "source": "sensor.energy",
+            },
+            "gas_meter": {
+                "source": "sensor.gas",
+            },
+        }
+    }
+
+    assert await async_setup_component(hass, DOMAIN, config)
+    await hass.async_block_till_done()
+
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
+    entity_id_energy = config[DOMAIN]["energy_meter"]["source"]
+    hass.states.async_set(
+        entity_id_energy, 2, {ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR}
+    )
+    entity_id_gas = config[DOMAIN]["gas_meter"]["source"]
+    hass.states.async_set(
+        entity_id_gas, 2, {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 == "0"
+    assert state.attributes.get(ATTR_DEVICE_CLASS) is None
+    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    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_MEASUREMENT
+    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_MEASUREMENT
+    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.attributes.get(ATTR_DEVICE_CLASS) is None
+    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "some_archaic_unit"
+
+
 async def test_restore_state(hass):
     """Test utility sensor restore state."""
     last_reset = "2020-12-21T00:00:00.013073+00:00"
@@ -193,7 +257,6 @@ async def test_restore_state(hass):
     )
 
     assert await async_setup_component(hass, DOMAIN, config)
-    assert await async_setup_component(hass, SENSOR_DOMAIN, config)
     await hass.async_block_till_done()
 
     # restore from cache
@@ -230,7 +293,6 @@ async def test_net_consumption(hass):
     }
 
     assert await async_setup_component(hass, DOMAIN, config)
-    assert await async_setup_component(hass, SENSOR_DOMAIN, config)
     await hass.async_block_till_done()
 
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
@@ -265,7 +327,6 @@ async def test_non_net_consumption(hass):
     }
 
     assert await async_setup_component(hass, DOMAIN, config)
-    assert await async_setup_component(hass, SENSOR_DOMAIN, config)
     await hass.async_block_till_done()
 
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
@@ -310,7 +371,6 @@ async def _test_self_reset(hass, config, start_time, expect_reset=True):
     now = dt_util.parse_datetime(start_time)
     with alter_time(now):
         assert await async_setup_component(hass, DOMAIN, config)
-        assert await async_setup_component(hass, SENSOR_DOMAIN, config)
         await hass.async_block_till_done()
 
         hass.bus.async_fire(EVENT_HOMEASSISTANT_START)