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

View File

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