diff --git a/custom_components/robovac/vacuum.py b/custom_components/robovac/vacuum.py
index 48ffafa..1b4c7cf 100644
--- a/custom_components/robovac/vacuum.py
+++ b/custom_components/robovac/vacuum.py
@@ -30,11 +30,7 @@
 from homeassistant.components.vacuum import (
     StateVacuumEntity,
     VacuumEntityFeature,
-    STATE_CLEANING,
-    STATE_DOCKED,
-    STATE_ERROR,
-    STATE_IDLE,
-    STATE_RETURNING,
+    VacuumActivity,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
@@ -208,15 +204,15 @@ def state(self) -> str | None:
                     getErrorMessage(self.error_code)
                 )
             )
-            return STATE_ERROR
+            return VacuumActivity.ERROR
         elif self.tuya_state == "Charging" or self.tuya_state == "completed":
-            return STATE_DOCKED
+            return VacuumActivity.DOCKED
         elif self.tuya_state == "Recharge":
-            return STATE_RETURNING
+            return VacuumActivity.RETURNING
         elif self.tuya_state == "Sleeping" or self.tuya_state == "standby":
-            return STATE_IDLE
+            return VacuumActivity.IDLE
         else:
-            return STATE_CLEANING
+            return VacuumActivity.CLEANING
 
     @property
     def extra_state_attributes(self) -> dict[str, Any]: