From a7ef983a3141b1e2147c318e72328c2ca08e123c Mon Sep 17 00:00:00 2001
From: Dave T <17680170+davet2001@users.noreply.github.com>
Date: Wed, 22 Dec 2021 09:00:24 +0000
Subject: [PATCH] Use new enums in litterrobot tests (#62515)

* Use new enums in litterrobot tests

* Code review: swap == for is
---
 tests/components/litterrobot/test_button.py | 10 +++-------
 tests/components/litterrobot/test_select.py |  5 +++--
 tests/components/litterrobot/test_sensor.py |  6 +++---
 tests/components/litterrobot/test_switch.py | 10 +++-------
 4 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/tests/components/litterrobot/test_button.py b/tests/components/litterrobot/test_button.py
index 0ca74da5d02..3f802d0e6b2 100644
--- a/tests/components/litterrobot/test_button.py
+++ b/tests/components/litterrobot/test_button.py
@@ -4,14 +4,10 @@ from unittest.mock import MagicMock
 from freezegun import freeze_time
 
 from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS
-from homeassistant.const import (
-    ATTR_ENTITY_ID,
-    ATTR_ICON,
-    ENTITY_CATEGORY_CONFIG,
-    STATE_UNKNOWN,
-)
+from homeassistant.const import ATTR_ENTITY_ID, ATTR_ICON, STATE_UNKNOWN
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers.entity import EntityCategory
 
 from .conftest import setup_integration
 
@@ -31,7 +27,7 @@ async def test_button(hass: HomeAssistant, mock_account: MagicMock) -> None:
 
     entry = entity_registry.async_get(BUTTON_ENTITY)
     assert entry
-    assert entry.entity_category == ENTITY_CATEGORY_CONFIG
+    assert entry.entity_category is EntityCategory.CONFIG
 
     await hass.services.async_call(
         BUTTON_DOMAIN,
diff --git a/tests/components/litterrobot/test_select.py b/tests/components/litterrobot/test_select.py
index dfb1b0e639e..e3e9782423e 100644
--- a/tests/components/litterrobot/test_select.py
+++ b/tests/components/litterrobot/test_select.py
@@ -10,9 +10,10 @@ from homeassistant.components.select import (
     DOMAIN as PLATFORM_DOMAIN,
     SERVICE_SELECT_OPTION,
 )
-from homeassistant.const import ATTR_ENTITY_ID, ENTITY_CATEGORY_CONFIG
+from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.util.dt import utcnow
 
 from .conftest import setup_integration
@@ -32,7 +33,7 @@ async def test_wait_time_select(hass: HomeAssistant, mock_account):
     ent_reg = entity_registry.async_get(hass)
     entity_entry = ent_reg.async_get(SELECT_ENTITY_ID)
     assert entity_entry
-    assert entity_entry.entity_category == ENTITY_CATEGORY_CONFIG
+    assert entity_entry.entity_category is EntityCategory.CONFIG
 
     data = {ATTR_ENTITY_ID: SELECT_ENTITY_ID}
 
diff --git a/tests/components/litterrobot/test_sensor.py b/tests/components/litterrobot/test_sensor.py
index dbc8c39790c..e3e62d5f5e4 100644
--- a/tests/components/litterrobot/test_sensor.py
+++ b/tests/components/litterrobot/test_sensor.py
@@ -2,8 +2,8 @@
 from unittest.mock import Mock
 
 from homeassistant.components.litterrobot.sensor import LitterRobotSleepTimeSensor
-from homeassistant.components.sensor import DOMAIN as PLATFORM_DOMAIN
-from homeassistant.const import DEVICE_CLASS_TIMESTAMP, PERCENTAGE
+from homeassistant.components.sensor import DOMAIN as PLATFORM_DOMAIN, SensorDeviceClass
+from homeassistant.const import PERCENTAGE
 
 from .conftest import create_mock_robot, setup_integration
 
@@ -30,7 +30,7 @@ async def test_sleep_time_sensor_with_none_state(hass):
 
     assert sensor
     assert sensor.state is None
-    assert sensor.device_class == DEVICE_CLASS_TIMESTAMP
+    assert sensor.device_class is SensorDeviceClass.TIMESTAMP
 
 
 async def test_gauge_icon():
diff --git a/tests/components/litterrobot/test_switch.py b/tests/components/litterrobot/test_switch.py
index 99c34e4273f..540a1c92810 100644
--- a/tests/components/litterrobot/test_switch.py
+++ b/tests/components/litterrobot/test_switch.py
@@ -10,14 +10,10 @@ from homeassistant.components.switch import (
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
 )
-from homeassistant.const import (
-    ATTR_ENTITY_ID,
-    ENTITY_CATEGORY_CONFIG,
-    STATE_OFF,
-    STATE_ON,
-)
+from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.util.dt import utcnow
 
 from .conftest import setup_integration
@@ -39,7 +35,7 @@ async def test_switch(hass: HomeAssistant, mock_account: MagicMock):
     ent_reg = entity_registry.async_get(hass)
     entity_entry = ent_reg.async_get(NIGHT_LIGHT_MODE_ENTITY_ID)
     assert entity_entry
-    assert entity_entry.entity_category == ENTITY_CATEGORY_CONFIG
+    assert entity_entry.entity_category is EntityCategory.CONFIG
 
 
 @pytest.mark.parametrize(