diff --git a/eufy_robovac/vacuum.py b/eufy_robovac/vacuum.py
index 0e328fd..c72b7b7 100644
--- a/eufy_robovac/vacuum.py
+++ b/eufy_robovac/vacuum.py
@@ -8,7 +8,7 @@
     SUPPORT_BATTERY, SUPPORT_CLEAN_SPOT, SUPPORT_FAN_SPEED, SUPPORT_LOCATE,
     SUPPORT_PAUSE, SUPPORT_RETURN_HOME, SUPPORT_STATUS, SUPPORT_START,
     SUPPORT_TURN_ON, SUPPORT_TURN_OFF,
-    VacuumEntity)
+    StateVacuumEntity)
 
 
 from . import robovac
@@ -50,7 +50,7 @@ def setup_platform(hass, config, add_entities, device_config=None):
     add_entities([EufyVacuum(device_config)], True)
 
 
-class EufyVacuum(VacuumEntity):
+class EufyVacuum(StateVacuumEntity):
     """Representation of a Eufy vacuum cleaner."""
 
     def __init__(self, device_config):