From bd7cc345ab3431ca0cc4b540cd1e6e2f6b3a89a5 Mon Sep 17 00:00:00 2001
From: G Johansson <goran.johansson@shiftit.se>
Date: Sun, 13 Oct 2024 14:30:12 +0000
Subject: [PATCH] Fix last test

---
 tests/components/litterrobot/test_vacuum.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/tests/components/litterrobot/test_vacuum.py b/tests/components/litterrobot/test_vacuum.py
index 735ee6653aa4cc..bf0d1569d0dc07 100644
--- a/tests/components/litterrobot/test_vacuum.py
+++ b/tests/components/litterrobot/test_vacuum.py
@@ -15,9 +15,7 @@
     DOMAIN as PLATFORM_DOMAIN,
     SERVICE_START,
     SERVICE_STOP,
-    STATE_DOCKED,
-    STATE_ERROR,
-    STATE_PAUSED,
+    VacuumEntityState,
 )
 from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
@@ -53,7 +51,7 @@ async def test_vacuum(
 
     vacuum = hass.states.get(VACUUM_ENTITY_ID)
     assert vacuum
-    assert vacuum.state == STATE_DOCKED
+    assert vacuum.state == VacuumEntityState.DOCKED
     assert vacuum.attributes["is_sleeping"] is False
 
     ent_reg_entry = entity_registry.async_get(VACUUM_ENTITY_ID)
@@ -95,15 +93,18 @@ async def test_vacuum_with_error(
 
     vacuum = hass.states.get(VACUUM_ENTITY_ID)
     assert vacuum
-    assert vacuum.state == STATE_ERROR
+    assert vacuum.state == VacuumEntityState.ERROR
 
 
 @pytest.mark.parametrize(
     ("robot_data", "expected_state"),
     [
-        ({"displayCode": "DC_CAT_DETECT"}, STATE_DOCKED),
-        ({"isDFIFull": True}, STATE_ERROR),
-        ({"robotCycleState": "CYCLE_STATE_CAT_DETECT"}, STATE_PAUSED),
+        ({"displayCode": "DC_CAT_DETECT"}, VacuumEntityState.DOCKED),
+        ({"isDFIFull": True}, VacuumEntityState.ERROR),
+        (
+            {"robotCycleState": "CYCLE_VacuumEntityState.CAT_DETECT"},
+            VacuumEntityState.PAUSED,
+        ),
     ],
 )
 async def test_vacuum_states(
@@ -150,7 +151,7 @@ async def test_commands(
 
     vacuum = hass.states.get(VACUUM_ENTITY_ID)
     assert vacuum
-    assert vacuum.state == STATE_DOCKED
+    assert vacuum.state == VacuumEntityState.DOCKED
 
     extra = extra or {}
     data = {ATTR_ENTITY_ID: VACUUM_ENTITY_ID, **extra.get("data", {})}