Use entity category enum in Rituals (home-assistant#60809)

iDVB · Dec 2, 2021 · 43d8c8f · 43d8c8f
1 parent 7dc2a11
commit 43d8c8f
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 14 deletions.
diff --git a/homeassistant/components/rituals_perfume_genie/binary_sensor.py b/homeassistant/components/rituals_perfume_genie/binary_sensor.py
@@ -8,8 +8,8 @@
     BinarySensorEntity,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import ENTITY_CATEGORY_DIAGNOSTIC
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import RitualsDataUpdateCoordinator
@@ -39,7 +39,7 @@ class DiffuserBatteryChargingBinarySensor(DiffuserEntity, BinarySensorEntity):
     """Representation of a diffuser battery charging binary sensor."""
 
     _attr_device_class = BinarySensorDeviceClass.BATTERY_CHARGING
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
 
     def __init__(
         self, diffuser: Diffuser, coordinator: RitualsDataUpdateCoordinator

diff --git a/homeassistant/components/rituals_perfume_genie/select.py b/homeassistant/components/rituals_perfume_genie/select.py
@@ -5,8 +5,9 @@
 
 from homeassistant.components.select import SelectEntity
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import AREA_SQUARE_METERS, ENTITY_CATEGORY_CONFIG
+from homeassistant.const import AREA_SQUARE_METERS
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import RitualsDataUpdateCoordinator
@@ -36,7 +37,7 @@ class DiffuserRoomSize(DiffuserEntity, SelectEntity):
     _attr_icon = "mdi:ruler-square"
     _attr_unit_of_measurement = AREA_SQUARE_METERS
     _attr_options = ["15", "30", "60", "100"]
-    _attr_entity_category = ENTITY_CATEGORY_CONFIG
+    _attr_entity_category = EntityCategory.CONFIG
 
     def __init__(
         self, diffuser: Diffuser, coordinator: RitualsDataUpdateCoordinator

diff --git a/homeassistant/components/rituals_perfume_genie/sensor.py b/homeassistant/components/rituals_perfume_genie/sensor.py
@@ -5,8 +5,9 @@
 
 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 EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import RitualsDataUpdateCoordinator
@@ -88,7 +89,7 @@ class DiffuserBatterySensor(DiffuserEntity, SensorEntity):
 
     _attr_device_class = SensorDeviceClass.BATTERY
     _attr_native_unit_of_measurement = PERCENTAGE
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
 
     def __init__(
         self, diffuser: Diffuser, coordinator: RitualsDataUpdateCoordinator
@@ -107,7 +108,7 @@ class DiffuserWifiSensor(DiffuserEntity, SensorEntity):
 
     _attr_device_class = SensorDeviceClass.SIGNAL_STRENGTH
     _attr_native_unit_of_measurement = PERCENTAGE
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
 
     def __init__(
         self, diffuser: Diffuser, coordinator: RitualsDataUpdateCoordinator

diff --git a/tests/components/rituals_perfume_genie/test_binary_sensor.py b/tests/components/rituals_perfume_genie/test_binary_sensor.py
@@ -1,9 +1,10 @@
 """Tests for the Rituals Perfume Genie binary sensor platform."""
 from homeassistant.components.binary_sensor import BinarySensorDeviceClass
 from homeassistant.components.rituals_perfume_genie.binary_sensor import CHARGING_SUFFIX
-from homeassistant.const import ATTR_DEVICE_CLASS, ENTITY_CATEGORY_DIAGNOSTIC, STATE_ON
+from homeassistant.const import ATTR_DEVICE_CLASS, STATE_ON
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
+from homeassistant.helpers.entity import EntityCategory
 
 from .common import (
     init_integration,
@@ -30,4 +31,4 @@ async def test_binary_sensors(hass: HomeAssistant) -> None:
     entry = registry.async_get("binary_sensor.genie_battery_charging")
     assert entry
     assert entry.unique_id == f"{hublot}{CHARGING_SUFFIX}"
-    assert entry.entity_category == ENTITY_CATEGORY_DIAGNOSTIC
+    assert entry.entity_category == EntityCategory.DIAGNOSTIC
diff --git a/tests/components/rituals_perfume_genie/test_select.py b/tests/components/rituals_perfume_genie/test_select.py
@@ -9,11 +9,11 @@
     AREA_SQUARE_METERS,
     ATTR_ENTITY_ID,
     ATTR_ICON,
-    ENTITY_CATEGORY_CONFIG,
     SERVICE_SELECT_OPTION,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.setup import async_setup_component
 
 from .common import init_integration, mock_config_entry, mock_diffuser
@@ -37,7 +37,7 @@ async def test_select_entity(hass: HomeAssistant) -> None:
     assert entry
     assert entry.unique_id == f"{diffuser.hublot}{ROOM_SIZE_SUFFIX}"
     assert entry.unit_of_measurement == AREA_SQUARE_METERS
-    assert entry.entity_category == ENTITY_CATEGORY_CONFIG
+    assert entry.entity_category == EntityCategory.CONFIG
 
 
 async def test_select_option(hass: HomeAssistant) -> None:

diff --git a/tests/components/rituals_perfume_genie/test_sensor.py b/tests/components/rituals_perfume_genie/test_sensor.py
@@ -10,11 +10,11 @@
     ATTR_DEVICE_CLASS,
     ATTR_ICON,
     ATTR_UNIT_OF_MEASUREMENT,
-    ENTITY_CATEGORY_DIAGNOSTIC,
     PERCENTAGE,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
+from homeassistant.helpers.entity import EntityCategory
 
 from .common import (
     init_integration,
@@ -59,7 +59,7 @@ async def test_sensors_diffuser_v1_battery_cartridge(hass: HomeAssistant) -> Non
     entry = registry.async_get("sensor.genie_battery")
     assert entry
     assert entry.unique_id == f"{hublot}{BATTERY_SUFFIX}"
-    assert entry.entity_category == ENTITY_CATEGORY_DIAGNOSTIC
+    assert entry.entity_category == EntityCategory.DIAGNOSTIC
 
     state = hass.states.get("sensor.genie_wifi")
     assert state
@@ -70,7 +70,7 @@ async def test_sensors_diffuser_v1_battery_cartridge(hass: HomeAssistant) -> Non
     entry = registry.async_get("sensor.genie_wifi")
     assert entry
     assert entry.unique_id == f"{hublot}{WIFI_SUFFIX}"
-    assert entry.entity_category == ENTITY_CATEGORY_DIAGNOSTIC
+    assert entry.entity_category == EntityCategory.DIAGNOSTIC
 
 
 async def test_sensors_diffuser_v2_no_battery_no_cartridge(hass: HomeAssistant) -> None: