From 25b1dfb53a464c28a26cb30f06c15e2521098d06 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Tue, 20 Sep 2022 16:49:44 +0200
Subject: [PATCH] Cleanup EntityCategory in tests (#78808)

---
 tests/components/hue/test_sensor_v1.py | 4 ++--
 tests/components/zha/test_button.py    | 7 +++----
 tests/components/zha/test_number.py    | 5 +++--
 tests/components/zha/test_select.py    | 9 +++++----
 4 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/tests/components/hue/test_sensor_v1.py b/tests/components/hue/test_sensor_v1.py
index c35403eaac1..acefbdd07fe 100644
--- a/tests/components/hue/test_sensor_v1.py
+++ b/tests/components/hue/test_sensor_v1.py
@@ -7,7 +7,7 @@ import aiohue
 from homeassistant.components import hue
 from homeassistant.components.hue.const import ATTR_HUE_EVENT
 from homeassistant.components.hue.v1 import sensor_base
-from homeassistant.const import ENTITY_CATEGORY_DIAGNOSTIC
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_registry import async_get
 from homeassistant.util import dt as dt_util
 
@@ -351,7 +351,7 @@ async def test_sensors(hass, mock_bridge_v1):
     ent_reg = async_get(hass)
     assert (
         ent_reg.async_get("sensor.hue_dimmer_switch_1_battery_level").entity_category
-        == ENTITY_CATEGORY_DIAGNOSTIC
+        == EntityCategory.DIAGNOSTIC
     )
 
 
diff --git a/tests/components/zha/test_button.py b/tests/components/zha/test_button.py
index 5e1350bc673..17c8f0ccb28 100644
--- a/tests/components/zha/test_button.py
+++ b/tests/components/zha/test_button.py
@@ -24,12 +24,11 @@ from homeassistant.components.button import DOMAIN, SERVICE_PRESS, ButtonDeviceC
 from homeassistant.const import (
     ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
-    ENTITY_CATEGORY_CONFIG,
-    ENTITY_CATEGORY_DIAGNOSTIC,
     STATE_UNKNOWN,
     Platform,
 )
 from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers.entity import EntityCategory
 
 from .common import find_entity_id
 from .conftest import SIG_EP_INPUT, SIG_EP_OUTPUT, SIG_EP_TYPE
@@ -137,7 +136,7 @@ async def test_button(hass, contact_sensor):
 
     entry = entity_registry.async_get(entity_id)
     assert entry
-    assert entry.entity_category == ENTITY_CATEGORY_DIAGNOSTIC
+    assert entry.entity_category == EntityCategory.DIAGNOSTIC
 
     with patch(
         "zigpy.zcl.Cluster.request",
@@ -177,7 +176,7 @@ async def test_frost_unlock(hass, tuya_water_valve):
 
     entry = entity_registry.async_get(entity_id)
     assert entry
-    assert entry.entity_category == ENTITY_CATEGORY_CONFIG
+    assert entry.entity_category == EntityCategory.CONFIG
 
     with patch(
         "zigpy.zcl.Cluster.request",
diff --git a/tests/components/zha/test_number.py b/tests/components/zha/test_number.py
index 808649b7f55..0bb620e98f4 100644
--- a/tests/components/zha/test_number.py
+++ b/tests/components/zha/test_number.py
@@ -8,8 +8,9 @@ import zigpy.zcl.clusters.general as general
 import zigpy.zcl.foundation as zcl_f
 
 from homeassistant.components.number import DOMAIN as NUMBER_DOMAIN
-from homeassistant.const import ENTITY_CATEGORY_CONFIG, STATE_UNAVAILABLE, Platform
+from homeassistant.const import STATE_UNAVAILABLE, Platform
 from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.setup import async_setup_component
 
 from .common import (
@@ -259,7 +260,7 @@ async def test_level_control_number(
 
     entity_entry = entity_registry.async_get(entity_id)
     assert entity_entry
-    assert entity_entry.entity_category == ENTITY_CATEGORY_CONFIG
+    assert entity_entry.entity_category == EntityCategory.CONFIG
 
     # Test number set_value
     await hass.services.async_call(
diff --git a/tests/components/zha/test_select.py b/tests/components/zha/test_select.py
index 1c714def54b..b9c72975823 100644
--- a/tests/components/zha/test_select.py
+++ b/tests/components/zha/test_select.py
@@ -8,8 +8,9 @@ import zigpy.profiles.zha as zha
 import zigpy.zcl.clusters.general as general
 import zigpy.zcl.clusters.security as security
 
-from homeassistant.const import ENTITY_CATEGORY_CONFIG, STATE_UNKNOWN, Platform
+from homeassistant.const import STATE_UNKNOWN, Platform
 from homeassistant.helpers import entity_registry as er, restore_state
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.util import dt as dt_util
 
 from .common import find_entity_id
@@ -136,7 +137,7 @@ async def test_select(hass, siren):
 
     entity_entry = entity_registry.async_get(entity_id)
     assert entity_entry
-    assert entity_entry.entity_category == ENTITY_CATEGORY_CONFIG
+    assert entity_entry.entity_category == EntityCategory.CONFIG
 
     # Test select option with string value
     await hass.services.async_call(
@@ -228,7 +229,7 @@ async def test_on_off_select_new_join(hass, light, zha_device_joined):
 
     entity_entry = entity_registry.async_get(entity_id)
     assert entity_entry
-    assert entity_entry.entity_category == ENTITY_CATEGORY_CONFIG
+    assert entity_entry.entity_category == EntityCategory.CONFIG
 
     # Test select option with string value
     await hass.services.async_call(
@@ -300,7 +301,7 @@ async def test_on_off_select_restored(hass, light, zha_device_restored):
 
     entity_entry = entity_registry.async_get(entity_id)
     assert entity_entry
-    assert entity_entry.entity_category == ENTITY_CATEGORY_CONFIG
+    assert entity_entry.entity_category == EntityCategory.CONFIG
 
 
 async def test_on_off_select_unsupported(hass, light, zha_device_joined_restored):