diff --git a/homeassistant/components/neato/sensor.py b/homeassistant/components/neato/sensor.py index 3f7b925ef7f..7cc2d0f171a 100644 --- a/homeassistant/components/neato/sensor.py +++ b/homeassistant/components/neato/sensor.py @@ -8,11 +8,11 @@ from typing import Any from pybotvac.exceptions import NeatoRobotException from pybotvac.robot import Robot -from homeassistant.components.sensor import DEVICE_CLASS_BATTERY, SensorEntity +from homeassistant.components.sensor import SensorDeviceClass, SensorEntity from homeassistant.config_entries import ConfigEntry -from homeassistant.const import ENTITY_CATEGORY_DIAGNOSTIC, PERCENTAGE +from homeassistant.const import PERCENTAGE from homeassistant.core import HomeAssistant -from homeassistant.helpers.entity import DeviceInfo +from homeassistant.helpers.entity import DeviceInfo, EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback from . import NeatoHub @@ -81,12 +81,12 @@ class NeatoSensor(SensorEntity): @property def device_class(self) -> str: """Return the device class.""" - return DEVICE_CLASS_BATTERY + return SensorDeviceClass.BATTERY @property def entity_category(self) -> str: """Device entity category.""" - return ENTITY_CATEGORY_DIAGNOSTIC + return EntityCategory.DIAGNOSTIC @property def available(self) -> bool: diff --git a/homeassistant/components/neato/switch.py b/homeassistant/components/neato/switch.py index c34eea492e9..f66765ecf0d 100644 --- a/homeassistant/components/neato/switch.py +++ b/homeassistant/components/neato/switch.py @@ -9,9 +9,9 @@ from pybotvac.exceptions import NeatoRobotException from pybotvac.robot import Robot from homeassistant.config_entries import ConfigEntry -from homeassistant.const import ENTITY_CATEGORY_CONFIG, STATE_OFF, STATE_ON +from homeassistant.const import STATE_OFF, STATE_ON from homeassistant.core import HomeAssistant -from homeassistant.helpers.entity import DeviceInfo, ToggleEntity +from homeassistant.helpers.entity import DeviceInfo, EntityCategory, ToggleEntity from homeassistant.helpers.entity_platform import AddEntitiesCallback from . import NeatoHub @@ -109,7 +109,7 @@ class NeatoConnectedSwitch(ToggleEntity): @property def entity_category(self) -> str: """Device entity category.""" - return ENTITY_CATEGORY_CONFIG + return EntityCategory.CONFIG @property def device_info(self) -> DeviceInfo: