Finish using enums in srp_energy (#62192)

This commit is contained in:
Robert Hillis 2021-12-19 07:02:29 -05:00 committed by GitHub
parent b01078199a
commit d52caf77d5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 15 deletions

View File

@ -5,12 +5,12 @@ import logging
import async_timeout
from requests.exceptions import ConnectionError as ConnectError, HTTPError, Timeout
from homeassistant.components.sensor import STATE_CLASS_TOTAL_INCREASING, SensorEntity
from homeassistant.const import (
ATTR_ATTRIBUTION,
DEVICE_CLASS_ENERGY,
ENERGY_KILO_WATT_HOUR,
from homeassistant.components.sensor import (
SensorDeviceClass,
SensorEntity,
SensorStateClass,
)
from homeassistant.const import ATTR_ATTRIBUTION, ENERGY_KILO_WATT_HOUR
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
from .const import (
@ -144,12 +144,12 @@ class SrpEntity(SensorEntity):
@property
def device_class(self):
"""Return the device class."""
return DEVICE_CLASS_ENERGY
return SensorDeviceClass.ENERGY
@property
def state_class(self):
"""Return the state class."""
return STATE_CLASS_TOTAL_INCREASING
return SensorStateClass.TOTAL_INCREASING
async def async_added_to_hass(self):
"""When entity is added to hass."""

View File

@ -1,7 +1,7 @@
"""Tests for the srp_energy sensor platform."""
from unittest.mock import MagicMock
from homeassistant.components.sensor import STATE_CLASS_TOTAL_INCREASING
from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
from homeassistant.components.srp_energy.const import (
ATTRIBUTION,
DEFAULT_NAME,
@ -11,11 +11,7 @@ from homeassistant.components.srp_energy.const import (
SRP_ENERGY_DOMAIN,
)
from homeassistant.components.srp_energy.sensor import SrpEntity, async_setup_entry
from homeassistant.const import (
ATTR_ATTRIBUTION,
DEVICE_CLASS_ENERGY,
ENERGY_KILO_WATT_HOUR,
)
from homeassistant.const import ATTR_ATTRIBUTION, ENERGY_KILO_WATT_HOUR
async def test_async_setup_entry(hass):
@ -99,8 +95,8 @@ async def test_srp_entity(hass):
assert srp_entity.should_poll is False
assert srp_entity.extra_state_attributes[ATTR_ATTRIBUTION] == ATTRIBUTION
assert srp_entity.available is not None
assert srp_entity.device_class == DEVICE_CLASS_ENERGY
assert srp_entity.state_class == STATE_CLASS_TOTAL_INCREASING
assert srp_entity.device_class is SensorDeviceClass.ENERGY
assert srp_entity.state_class is SensorStateClass.TOTAL_INCREASING
await srp_entity.async_added_to_hass()
assert srp_entity.state is not None