diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py
index be5b2d03e1c..deae3be3485 100644
--- a/homeassistant/components/unifi/sensor.py
+++ b/homeassistant/components/unifi/sensor.py
@@ -6,10 +6,11 @@ Support for uptime sensors of network clients.
 
 from datetime import datetime, timedelta
 
-from homeassistant.components.sensor import DEVICE_CLASS_TIMESTAMP, DOMAIN, SensorEntity
-from homeassistant.const import DATA_MEGABYTES, ENTITY_CATEGORY_DIAGNOSTIC
+from homeassistant.components.sensor import DOMAIN, SensorDeviceClass, SensorEntity
+from homeassistant.const import DATA_MEGABYTES
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.entity import EntityCategory
 import homeassistant.util.dt as dt_util
 
 from .const import DOMAIN as UNIFI_DOMAIN
@@ -86,7 +87,7 @@ class UniFiBandwidthSensor(UniFiClient, SensorEntity):
 
     DOMAIN = DOMAIN
 
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
     _attr_native_unit_of_measurement = DATA_MEGABYTES
 
     @property
@@ -132,8 +133,8 @@ class UniFiUpTimeSensor(UniFiClient, SensorEntity):
     DOMAIN = DOMAIN
     TYPE = UPTIME_SENSOR
 
-    _attr_device_class = DEVICE_CLASS_TIMESTAMP
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
+    _attr_device_class = SensorDeviceClass.TIMESTAMP
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
 
     def __init__(self, client, controller):
         """Set up tracked client."""
diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py
index 075718b4da5..638ff0aa397 100644
--- a/homeassistant/components/unifi/switch.py
+++ b/homeassistant/components/unifi/switch.py
@@ -15,11 +15,10 @@ from aiounifi.events import (
 )
 
 from homeassistant.components.switch import DOMAIN, SwitchEntity
-from homeassistant.const import ENTITY_CATEGORY_CONFIG
 from homeassistant.core import callback
 from homeassistant.helpers.device_registry import DeviceEntryType
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
-from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.entity import DeviceInfo, EntityCategory
 from homeassistant.helpers.entity_registry import async_entries_for_config_entry
 from homeassistant.helpers.restore_state import RestoreEntity
 
@@ -185,7 +184,7 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchEntity, RestoreEntity):
     DOMAIN = DOMAIN
     TYPE = POE_SWITCH
 
-    _attr_entity_category = ENTITY_CATEGORY_CONFIG
+    _attr_entity_category = EntityCategory.CONFIG
 
     def __init__(self, client, controller):
         """Set up POE switch."""
@@ -274,7 +273,7 @@ class UniFiBlockClientSwitch(UniFiClient, SwitchEntity):
     DOMAIN = DOMAIN
     TYPE = BLOCK_SWITCH
 
-    _attr_entity_category = ENTITY_CATEGORY_CONFIG
+    _attr_entity_category = EntityCategory.CONFIG
 
     def __init__(self, client, controller):
         """Set up block switch."""
@@ -325,7 +324,7 @@ class UniFiDPIRestrictionSwitch(UniFiBase, SwitchEntity):
     DOMAIN = DOMAIN
     TYPE = DPI_SWITCH
 
-    _attr_entity_category = ENTITY_CATEGORY_CONFIG
+    _attr_entity_category = EntityCategory.CONFIG
 
     @property
     def key(self) -> Any: