diff --git a/homeassistant/components/roomba/sensor.py b/homeassistant/components/roomba/sensor.py
index 4552159677c..ae33930b549 100644
--- a/homeassistant/components/roomba/sensor.py
+++ b/homeassistant/components/roomba/sensor.py
@@ -1,11 +1,8 @@
 """Sensor for checking the battery level of Roomba."""
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
 from homeassistant.components.vacuum import STATE_DOCKED
-from homeassistant.const import (
-    DEVICE_CLASS_BATTERY,
-    ENTITY_CATEGORY_DIAGNOSTIC,
-    PERCENTAGE,
-)
+from homeassistant.const import PERCENTAGE
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.icon import icon_for_battery_level
 
 from .const import BLID, DOMAIN, ROOMBA_SESSION
@@ -24,7 +21,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 class RoombaBattery(IRobotEntity, SensorEntity):
     """Class to hold Roomba Sensor basic info."""
 
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
 
     @property
     def name(self):
@@ -39,7 +36,7 @@ class RoombaBattery(IRobotEntity, SensorEntity):
     @property
     def device_class(self):
         """Return the device class of the sensor."""
-        return DEVICE_CLASS_BATTERY
+        return SensorDeviceClass.BATTERY
 
     @property
     def native_unit_of_measurement(self):