diff --git a/custom_components/dreame_vacuum/vacuum.py b/custom_components/dreame_vacuum/vacuum.py
index 96a1d27..23db90b 100644
--- a/custom_components/dreame_vacuum/vacuum.py
+++ b/custom_components/dreame_vacuum/vacuum.py
@@ -13,12 +13,7 @@
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.icon import icon_for_battery_level
 from homeassistant.components.vacuum import (
-    STATE_CLEANING,
-    STATE_DOCKED,
-    STATE_ERROR,
-    STATE_IDLE,
-    STATE_PAUSED,
-    STATE_RETURNING,
+    VacuumActivity,
     StateVacuumEntity,
     VacuumEntityFeature,
 )
@@ -133,43 +128,43 @@
 
 STATE_CODE_TO_STATE: Final = {
     DreameVacuumState.UNKNOWN: STATE_UNKNOWN,
-    DreameVacuumState.SWEEPING: STATE_CLEANING,
-    DreameVacuumState.IDLE: STATE_IDLE,
-    DreameVacuumState.PAUSED: STATE_PAUSED,
-    DreameVacuumState.ERROR: STATE_ERROR,
-    DreameVacuumState.RETURNING: STATE_RETURNING,
-    DreameVacuumState.CHARGING: STATE_DOCKED,
-    DreameVacuumState.MOPPING: STATE_CLEANING,
-    DreameVacuumState.DRYING: STATE_DOCKED,
-    DreameVacuumState.WASHING: STATE_CLEANING,
-    DreameVacuumState.RETURNING_TO_WASH: STATE_RETURNING,
-    DreameVacuumState.BUILDING: STATE_DOCKED,
-    DreameVacuumState.SWEEPING_AND_MOPPING: STATE_CLEANING,
-    DreameVacuumState.CHARGING_COMPLETED: STATE_DOCKED,
-    DreameVacuumState.UPGRADING: STATE_IDLE,
-    DreameVacuumState.CLEAN_SUMMON: STATE_CLEANING,
-    DreameVacuumState.STATION_RESET: STATE_IDLE,
-    DreameVacuumState.RETURNING_INSTALL_MOP: STATE_RETURNING,
-    DreameVacuumState.RETURNING_REMOVE_MOP: STATE_RETURNING,
-    DreameVacuumState.WATER_CHECK: STATE_DOCKED,
-    DreameVacuumState.CLEAN_ADD_WATER: STATE_CLEANING,
-    DreameVacuumState.WASHING_PAUSED: STATE_PAUSED,
-    DreameVacuumState.AUTO_EMPTYING: STATE_DOCKED,
-    DreameVacuumState.REMOTE_CONTROL: STATE_CLEANING,
-    DreameVacuumState.SMART_CHARGING: STATE_DOCKED,
-    DreameVacuumState.SECOND_CLEANING: STATE_CLEANING,
-    DreameVacuumState.HUMAN_FOLLOWING: STATE_CLEANING,
-    DreameVacuumState.SPOT_CLEANING: STATE_CLEANING,
-    DreameVacuumState.RETURNING_AUTO_EMPTY: STATE_RETURNING,
-    DreameVacuumState.SHORTCUT: STATE_CLEANING,
-    DreameVacuumState.WAITING_FOR_TASK: STATE_IDLE,
-    DreameVacuumState.STATION_CLEANING: STATE_CLEANING,
-    DreameVacuumState.RETURNING_TO_DRAIN: STATE_RETURNING,
-    DreameVacuumState.DRAINING: STATE_CLEANING,
-    DreameVacuumState.AUTO_WATER_DRAINING: STATE_CLEANING,
-    DreameVacuumState.SHORTCUT: STATE_CLEANING,
-    DreameVacuumState.MONITORING: STATE_CLEANING,
-    DreameVacuumState.MONITORING_PAUSED: STATE_PAUSED,
+    DreameVacuumState.SWEEPING: VacuumActivity.CLEANING,
+    DreameVacuumState.IDLE: VacuumActivity.IDLE,
+    DreameVacuumState.PAUSED: VacuumActivity.PAUSED,
+    DreameVacuumState.ERROR: VacuumActivity.ERROR,
+    DreameVacuumState.RETURNING: VacuumActivity.RETURNING,
+    DreameVacuumState.CHARGING: VacuumActivity.DOCKED,
+    DreameVacuumState.MOPPING: VacuumActivity.CLEANING,
+    DreameVacuumState.DRYING: VacuumActivity.DOCKED,
+    DreameVacuumState.WASHING: VacuumActivity.CLEANING,
+    DreameVacuumState.RETURNING_TO_WASH: VacuumActivity.RETURNING,
+    DreameVacuumState.BUILDING: VacuumActivity.DOCKED,
+    DreameVacuumState.SWEEPING_AND_MOPPING: VacuumActivity.CLEANING,
+    DreameVacuumState.CHARGING_COMPLETED: VacuumActivity.DOCKED,
+    DreameVacuumState.UPGRADING: VacuumActivity.IDLE,
+    DreameVacuumState.CLEAN_SUMMON: VacuumActivity.CLEANING,
+    DreameVacuumState.STATION_RESET: VacuumActivity.IDLE,
+    DreameVacuumState.RETURNING_INSTALL_MOP: VacuumActivity.RETURNING,
+    DreameVacuumState.RETURNING_REMOVE_MOP: VacuumActivity.RETURNING,
+    DreameVacuumState.WATER_CHECK: VacuumActivity.DOCKED,
+    DreameVacuumState.CLEAN_ADD_WATER: VacuumActivity.CLEANING,
+    DreameVacuumState.WASHING_PAUSED: VacuumActivity.PAUSED,
+    DreameVacuumState.AUTO_EMPTYING: VacuumActivity.DOCKED,
+    DreameVacuumState.REMOTE_CONTROL: VacuumActivity.CLEANING,
+    DreameVacuumState.SMART_CHARGING: VacuumActivity.DOCKED,
+    DreameVacuumState.SECOND_CLEANING: VacuumActivity.CLEANING,
+    DreameVacuumState.HUMAN_FOLLOWING: VacuumActivity.CLEANING,
+    DreameVacuumState.SPOT_CLEANING: VacuumActivity.CLEANING,
+    DreameVacuumState.RETURNING_AUTO_EMPTY: VacuumActivity.RETURNING,
+    DreameVacuumState.SHORTCUT: VacuumActivity.CLEANING,
+    DreameVacuumState.WAITING_FOR_TASK: VacuumActivity.IDLE,
+    DreameVacuumState.STATION_CLEANING: VacuumActivity.CLEANING,
+    DreameVacuumState.RETURNING_TO_DRAIN: VacuumActivity.RETURNING,
+    DreameVacuumState.DRAINING: VacuumActivity.CLEANING,
+    DreameVacuumState.AUTO_WATER_DRAINING: VacuumActivity.CLEANING,
+    DreameVacuumState.SHORTCUT: VacuumActivity.CLEANING,
+    DreameVacuumState.MONITORING: VacuumActivity.CLEANING,
+    DreameVacuumState.MONITORING_PAUSED: VacuumActivity.PAUSED,
 }
 
 SUCTION_LEVEL_TO_FAN_SPEED: Final = {