diff --git a/homeassistant/components/utility_meter/sensor.py b/homeassistant/components/utility_meter/sensor.py index 1d244c970ff..94a9e0d9175 100644 --- a/homeassistant/components/utility_meter/sensor.py +++ b/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.""" diff --git a/tests/components/utility_meter/test_sensor.py b/tests/components/utility_meter/test_sensor.py index 54854ac9668..c5075aa322b 100644 --- a/tests/components/utility_meter/test_sensor.py +++ b/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)