Integration. Add device class, last_reset, state_class (#53698)

Co-authored-by: Franck Nijhof <git@frenck.dev>
This commit is contained in:
Daniel Hjelseth Høyer 2021-07-29 20:05:53 +02:00 committed by GitHub
parent bedb9550f5
commit 6590e464af
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 132 additions and 13 deletions

View file

@ -4,8 +4,16 @@ import logging
import voluptuous as vol
from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
from homeassistant.components.sensor import (
ATTR_LAST_RESET,
DEVICE_CLASS_ENERGY,
DEVICE_CLASS_POWER,
PLATFORM_SCHEMA,
STATE_CLASS_MEASUREMENT,
SensorEntity,
)
from homeassistant.const import (
ATTR_DEVICE_CLASS,
ATTR_UNIT_OF_MEASUREMENT,
CONF_METHOD,
CONF_NAME,
@ -20,6 +28,7 @@ from homeassistant.core import callback
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.event import async_track_state_change_event
from homeassistant.helpers.restore_state import RestoreEntity
from homeassistant.util import dt as dt_util
# mypy: allow-untyped-defs, no-check-untyped-defs
@ -115,16 +124,26 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
self._unit_prefix = UNIT_PREFIXES[unit_prefix]
self._unit_time = UNIT_TIME[unit_time]
self._attr_state_class = STATE_CLASS_MEASUREMENT
async def async_added_to_hass(self):
"""Handle entity which will be added."""
await super().async_added_to_hass()
state = await self.async_get_last_state()
self._attr_last_reset = dt_util.utcnow()
if state:
try:
self._state = Decimal(state.state)
except ValueError as err:
except (DecimalException, ValueError) as err:
_LOGGER.warning("Could not restore last state: %s", err)
else:
last_reset = dt_util.parse_datetime(
state.attributes.get(ATTR_LAST_RESET, "")
)
self._attr_last_reset = (
last_reset if last_reset else dt_util.utc_from_timestamp(0)
)
self._attr_device_class = state.attributes.get(ATTR_DEVICE_CLASS)
@callback
def calc_integration(event):
@ -143,7 +162,11 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
self._unit_of_measurement = self._unit_template.format(
"" if unit is None else unit
)
if (
self.device_class is None
and new_state.attributes.get(ATTR_DEVICE_CLASS) == DEVICE_CLASS_POWER
):
self._attr_device_class = DEVICE_CLASS_ENERGY
try:
# integration as the Riemann integral of previous measures.
area = 0

View file

@ -2,12 +2,22 @@
from datetime import timedelta
from unittest.mock import patch
from homeassistant.const import ENERGY_KILO_WATT_HOUR, POWER_WATT, TIME_SECONDS
from homeassistant.components.sensor import STATE_CLASS_MEASUREMENT
from homeassistant.const import (
DEVICE_CLASS_ENERGY,
DEVICE_CLASS_POWER,
ENERGY_KILO_WATT_HOUR,
POWER_WATT,
TIME_SECONDS,
)
from homeassistant.core import HomeAssistant, State
from homeassistant.setup import async_setup_component
import homeassistant.util.dt as dt_util
from tests.common import mock_restore_cache
async def test_state(hass):
async def test_state(hass) -> None:
"""Test integration sensor state."""
config = {
"sensor": {
@ -19,15 +29,25 @@ async def test_state(hass):
}
}
now = dt_util.utcnow()
with patch("homeassistant.util.dt.utcnow", return_value=now):
assert await async_setup_component(hass, "sensor", config)
entity_id = config["sensor"]["source"]
hass.states.async_set(entity_id, 1, {})
await hass.async_block_till_done()
now = dt_util.utcnow() + timedelta(seconds=3600)
with patch("homeassistant.util.dt.utcnow", return_value=now):
hass.states.async_set(entity_id, 1, {}, force_update=True)
state = hass.states.get("sensor.integration")
assert state is not None
assert state.attributes.get("last_reset") == now.isoformat()
assert state.attributes.get("state_class") == STATE_CLASS_MEASUREMENT
assert "device_class" not in state.attributes
future_now = dt_util.utcnow() + timedelta(seconds=3600)
with patch("homeassistant.util.dt.utcnow", return_value=future_now):
hass.states.async_set(
entity_id, 1, {"device_class": DEVICE_CLASS_POWER}, force_update=True
)
await hass.async_block_till_done()
state = hass.states.get("sensor.integration")
@ -37,6 +57,82 @@ async def test_state(hass):
assert round(float(state.state), config["sensor"]["round"]) == 1.0
assert state.attributes.get("unit_of_measurement") == ENERGY_KILO_WATT_HOUR
assert state.attributes.get("device_class") == DEVICE_CLASS_ENERGY
assert state.attributes.get("state_class") == STATE_CLASS_MEASUREMENT
assert state.attributes.get("last_reset") == now.isoformat()
async def test_restore_state(hass: HomeAssistant) -> None:
"""Test integration sensor state is restored correctly."""
mock_restore_cache(
hass,
(
State(
"sensor.integration",
"100.0",
{
"last_reset": "2019-10-06T21:00:00",
"device_class": DEVICE_CLASS_ENERGY,
},
),
),
)
config = {
"sensor": {
"platform": "integration",
"name": "integration",
"source": "sensor.power",
"unit": ENERGY_KILO_WATT_HOUR,
"round": 2,
}
}
assert await async_setup_component(hass, "sensor", config)
await hass.async_block_till_done()
state = hass.states.get("sensor.integration")
assert state
assert state.state == "100.00"
assert state.attributes.get("unit_of_measurement") == ENERGY_KILO_WATT_HOUR
assert state.attributes.get("device_class") == DEVICE_CLASS_ENERGY
assert state.attributes.get("last_reset") == "2019-10-06T21:00:00"
async def test_restore_state_failed(hass: HomeAssistant) -> None:
"""Test integration sensor state is restored correctly."""
mock_restore_cache(
hass,
(
State(
"sensor.integration",
"INVALID",
{
"last_reset": "2019-10-06T21:00:00.000000",
},
),
),
)
config = {
"sensor": {
"platform": "integration",
"name": "integration",
"source": "sensor.power",
"unit": ENERGY_KILO_WATT_HOUR,
}
}
assert await async_setup_component(hass, "sensor", config)
await hass.async_block_till_done()
state = hass.states.get("sensor.integration")
assert state
assert state.state == "0"
assert state.attributes.get("unit_of_measurement") == ENERGY_KILO_WATT_HOUR
assert state.attributes.get("state_class") == STATE_CLASS_MEASUREMENT
assert state.attributes.get("last_reset") != "2019-10-06T21:00:00"
assert "device_class" not in state.attributes
async def test_trapezoidal(hass):