From 85c122f3e560e088221d8194904d3952b3f88e36 Mon Sep 17 00:00:00 2001
From: Robert Hillis <tkdrob4390@yahoo.com>
Date: Fri, 17 Dec 2021 05:12:43 -0500
Subject: [PATCH] Use enums in unifi tests (#62151)

---
 tests/components/unifi/test_sensor.py | 6 +++---
 tests/components/unifi/test_switch.py | 5 +++--
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py
index 3794c46988d..c0b45efb119 100644
--- a/tests/components/unifi/test_sensor.py
+++ b/tests/components/unifi/test_sensor.py
@@ -15,9 +15,9 @@ from homeassistant.components.unifi.const import (
     CONF_TRACK_DEVICES,
     DOMAIN as UNIFI_DOMAIN,
 )
-from homeassistant.const import ENTITY_CATEGORY_DIAGNOSTIC
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.entity import EntityCategory
 import homeassistant.util.dt as dt_util
 
 from .test_controller import setup_unifi_integration
@@ -79,7 +79,7 @@ async def test_bandwidth_sensors(hass, aioclient_mock, mock_unifi_websocket):
     ent_reg = er.async_get(hass)
     assert (
         ent_reg.async_get("sensor.wired_client_rx").entity_category
-        == ENTITY_CATEGORY_DIAGNOSTIC
+        is EntityCategory.DIAGNOSTIC
     )
 
     # Verify state update
@@ -190,7 +190,7 @@ async def test_uptime_sensors(
     ent_reg = er.async_get(hass)
     assert (
         ent_reg.async_get("sensor.client1_uptime").entity_category
-        == ENTITY_CATEGORY_DIAGNOSTIC
+        is EntityCategory.DIAGNOSTIC
     )
 
     # Verify normal new event doesn't change uptime
diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py
index 29640b7a4b4..38231ef9609 100644
--- a/tests/components/unifi/test_switch.py
+++ b/tests/components/unifi/test_switch.py
@@ -16,9 +16,10 @@ from homeassistant.components.unifi.const import (
     DOMAIN as UNIFI_DOMAIN,
 )
 from homeassistant.components.unifi.switch import POE_SWITCH
-from homeassistant.const import ENTITY_CATEGORY_CONFIG, STATE_OFF, STATE_ON
+from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.entity import EntityCategory
 
 from .test_controller import (
     CONTROLLER_HOST,
@@ -408,7 +409,7 @@ async def test_switches(hass, aioclient_mock):
         "switch.block_client_1",
         "switch.block_media_streaming",
     ):
-        assert ent_reg.async_get(entry_id).entity_category == ENTITY_CATEGORY_CONFIG
+        assert ent_reg.async_get(entry_id).entity_category is EntityCategory.CONFIG
 
     # Block and unblock client