Add support for STATES of vacuums (#15573)

* Vacuum: Added support for STATES

* Added debug logging and corrected state order

* typo

* Fix travis error, STATE = STATE for readability

* status -> state

* Changed to Entity instead of ToogleEntity

* Updated some vacuums

* Revert changes

* Revert Changes

* added SUPPORT_STATE

* Woof?

* Implement on/off if STATE not supported

* Moved new state vaccum to Class StateVacuumDevice

* Error: I should go to bed

* Moved around methods for easier reading

* Added StateVacuumDevice demo vacuum

* Added tests for StateVacuumDevice demo vacuum

* Fix styling errors

* Refactored to BaseVaccum

* Vacuum will now go back to dock

* Class BaseVacuum is for internal use only

* return -> await

* return -> await

homeassistant/components/vacuum/__init__.py

@@ -14,12 +14,12 @@ import voluptuous as vol
 from homeassistant.components import group
 from homeassistant.const import (
     ATTR_BATTERY_LEVEL, ATTR_COMMAND, ATTR_ENTITY_ID, SERVICE_TOGGLE,
-    SERVICE_TURN_OFF, SERVICE_TURN_ON, STATE_ON)
+    SERVICE_TURN_OFF, SERVICE_TURN_ON, STATE_ON, STATE_PAUSED, STATE_IDLE)
 from homeassistant.loader import bind_hass
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.config_validation import PLATFORM_SCHEMA  # noqa
 from homeassistant.helpers.entity_component import EntityComponent
-from homeassistant.helpers.entity import ToggleEntity
+from homeassistant.helpers.entity import (ToggleEntity, Entity)
 from homeassistant.helpers.icon import icon_for_battery_level
 
 _LOGGER = logging.getLogger(__name__)
@@ -75,6 +75,13 @@ SERVICE_TO_METHOD = {
                            'schema': VACUUM_SEND_COMMAND_SERVICE_SCHEMA},
 }
 
+STATE_CLEANING = 'cleaning'
+STATE_DOCKED = 'docked'
+STATE_IDLE = STATE_IDLE
+STATE_PAUSED = STATE_PAUSED
+STATE_RETURNING = 'returning'
+STATE_ERROR = 'error'
+
 DEFAULT_NAME = 'Vacuum cleaner robot'
 
 SUPPORT_TURN_ON = 1
@@ -89,6 +96,7 @@ SUPPORT_SEND_COMMAND = 256
 SUPPORT_LOCATE = 512
 SUPPORT_CLEAN_SPOT = 1024
 SUPPORT_MAP = 2048
+SUPPORT_STATE = 4096
 
 
 @bind_hass
@@ -208,33 +216,22 @@ def async_setup(hass, config):
     return True
 
 
-class VacuumDevice(ToggleEntity):
+class _BaseVacuum(Entity):
-    """Representation of a vacuum cleaner robot."""
+    """Representation of a base vacuum.
+
+    Contains common properties and functions for all vacuum devices.
+    """
 
     @property
     def supported_features(self):
         """Flag vacuum cleaner features that are supported."""
         raise NotImplementedError()
 
-    @property
-    def status(self):
-        """Return the status of the vacuum cleaner."""
-        return None
-
     @property
     def battery_level(self):
         """Return the battery level of the vacuum cleaner."""
         return None
 
-    @property
-    def battery_icon(self):
-        """Return the battery icon for the vacuum cleaner."""
-        charging = False
-        if self.status is not None:
-            charging = 'charg' in self.status.lower()
-        return icon_for_battery_level(
-            battery_level=self.battery_level, charging=charging)
-
     @property
     def fan_speed(self):
         """Return the fan speed of the vacuum cleaner."""
@@ -245,6 +242,106 @@ class VacuumDevice(ToggleEntity):
         """Get the list of available fan speed steps of the vacuum cleaner."""
         raise NotImplementedError()
 
+    def start_pause(self, **kwargs):
+        """Start, pause or resume the cleaning task."""
+        raise NotImplementedError()
+
+    async def async_start_pause(self, **kwargs):
+        """Start, pause or resume the cleaning task.
+
+        This method must be run in the event loop.
+        """
+        await self.hass.async_add_executor_job(
+            partial(self.start_pause, **kwargs))
+
+    def stop(self, **kwargs):
+        """Stop the vacuum cleaner."""
+        raise NotImplementedError()
+
+    async def async_stop(self, **kwargs):
+        """Stop the vacuum cleaner.
+
+        This method must be run in the event loop.
+        """
+        await self.hass.async_add_executor_job(partial(self.stop, **kwargs))
+
+    def return_to_base(self, **kwargs):
+        """Set the vacuum cleaner to return to the dock."""
+        raise NotImplementedError()
+
+    async def async_return_to_base(self, **kwargs):
+        """Set the vacuum cleaner to return to the dock.
+
+        This method must be run in the event loop.
+        """
+        await self.hass.async_add_executor_job(
+            partial(self.return_to_base, **kwargs))
+
+    def clean_spot(self, **kwargs):
+        """Perform a spot clean-up."""
+        raise NotImplementedError()
+
+    async def async_clean_spot(self, **kwargs):
+        """Perform a spot clean-up.
+
+        This method must be run in the event loop.
+        """
+        await self.hass.async_add_executor_job(
+            partial(self.clean_spot, **kwargs))
+
+    def locate(self, **kwargs):
+        """Locate the vacuum cleaner."""
+        raise NotImplementedError()
+
+    async def async_locate(self, **kwargs):
+        """Locate the vacuum cleaner.
+
+        This method must be run in the event loop.
+        """
+        await self.hass.async_add_executor_job(partial(self.locate, **kwargs))
+
+    def set_fan_speed(self, fan_speed, **kwargs):
+        """Set fan speed."""
+        raise NotImplementedError()
+
+    async def async_set_fan_speed(self, fan_speed, **kwargs):
+        """Set fan speed.
+
+        This method must be run in the event loop.
+        """
+        await self.hass.async_add_executor_job(
+            partial(self.set_fan_speed, fan_speed, **kwargs))
+
+    def send_command(self, command, params=None, **kwargs):
+        """Send a command to a vacuum cleaner."""
+        raise NotImplementedError()
+
+    async def async_send_command(self, command, params=None, **kwargs):
+        """Send a command to a vacuum cleaner.
+
+        This method must be run in the event loop.
+        """
+        await self.hass.async_add_executor_job(
+            partial(self.send_command, command, params=params, **kwargs))
+
+
+class VacuumDevice(_BaseVacuum, ToggleEntity):
+    """Representation of a vacuum cleaner robot."""
+
+    @property
+    def status(self):
+        """Return the status of the vacuum cleaner."""
+        return None
+
+    @property
+    def battery_icon(self):
+        """Return the battery icon for the vacuum cleaner."""
+        charging = False
+        if self.status is not None:
+            charging = 'charg' in self.status.lower()
+        return icon_for_battery_level(
+            battery_level=self.battery_level, charging=charging)
+
     @property
     def state_attributes(self):
         """Return the state attributes of the vacuum cleaner."""
@@ -267,100 +364,54 @@ class VacuumDevice(ToggleEntity):
         """Turn the vacuum on and start cleaning."""
         raise NotImplementedError()
 
-    def async_turn_on(self, **kwargs):
+    async def async_turn_on(self, **kwargs):
         """Turn the vacuum on and start cleaning.
 
-        This method must be run in the event loop and returns a coroutine.
+        This method must be run in the event loop.
         """
-        return self.hass.async_add_job(partial(self.turn_on, **kwargs))
+        await self.hass.async_add_executor_job(
+            partial(self.turn_on, **kwargs))
 
     def turn_off(self, **kwargs):
         """Turn the vacuum off stopping the cleaning and returning home."""
         raise NotImplementedError()
 
-    def async_turn_off(self, **kwargs):
+    async def async_turn_off(self, **kwargs):
         """Turn the vacuum off stopping the cleaning and returning home.
 
-        This method must be run in the event loop and returns a coroutine.
+        This method must be run in the event loop.
         """
-        return self.hass.async_add_job(partial(self.turn_off, **kwargs))
+        await self.hass.async_add_executor_job(
+            partial(self.turn_off, **kwargs))
 
-    def return_to_base(self, **kwargs):
-        """Set the vacuum cleaner to return to the dock."""
-        raise NotImplementedError()
 
-    def async_return_to_base(self, **kwargs):
+class StateVacuumDevice(_BaseVacuum):
-        """Set the vacuum cleaner to return to the dock.
+    """Representation of a vacuum cleaner robot that supports states."""
 
-        This method must be run in the event loop and returns a coroutine.
+    @property
-        """
+    def state(self):
-        return self.hass.async_add_job(partial(self.return_to_base, **kwargs))
+        """Return the state of the vacuum cleaner."""
+        return None
 
-    def stop(self, **kwargs):
+    @property
-        """Stop the vacuum cleaner."""
+    def battery_icon(self):
-        raise NotImplementedError()
+        """Return the battery icon for the vacuum cleaner."""
+        charging = bool(self.state == STATE_DOCKED)
 
-    def async_stop(self, **kwargs):
+        return icon_for_battery_level(
-        """Stop the vacuum cleaner.
+            battery_level=self.battery_level, charging=charging)
 
-        This method must be run in the event loop and returns a coroutine.
+    @property
-        """
+    def state_attributes(self):
-        return self.hass.async_add_job(partial(self.stop, **kwargs))
+        """Return the state attributes of the vacuum cleaner."""
+        data = {}
 
-    def clean_spot(self, **kwargs):
+        if self.battery_level is not None:
-        """Perform a spot clean-up."""
+            data[ATTR_BATTERY_LEVEL] = self.battery_level
-        raise NotImplementedError()
+            data[ATTR_BATTERY_ICON] = self.battery_icon
 
-    def async_clean_spot(self, **kwargs):
+        if self.fan_speed is not None:
-        """Perform a spot clean-up.
+            data[ATTR_FAN_SPEED] = self.fan_speed
+            data[ATTR_FAN_SPEED_LIST] = self.fan_speed_list
 
-        This method must be run in the event loop and returns a coroutine.
+        return data
-        """
-        return self.hass.async_add_job(partial(self.clean_spot, **kwargs))
-
-    def locate(self, **kwargs):
-        """Locate the vacuum cleaner."""
-        raise NotImplementedError()
-
-    def async_locate(self, **kwargs):
-        """Locate the vacuum cleaner.
-
-        This method must be run in the event loop and returns a coroutine.
-        """
-        return self.hass.async_add_job(partial(self.locate, **kwargs))
-
-    def set_fan_speed(self, fan_speed, **kwargs):
-        """Set fan speed."""
-        raise NotImplementedError()
-
-    def async_set_fan_speed(self, fan_speed, **kwargs):
-        """Set fan speed.
-
-        This method must be run in the event loop and returns a coroutine.
-        """
-        return self.hass.async_add_job(
-            partial(self.set_fan_speed, fan_speed, **kwargs))
-
-    def start_pause(self, **kwargs):
-        """Start, pause or resume the cleaning task."""
-        raise NotImplementedError()
-
-    def async_start_pause(self, **kwargs):
-        """Start, pause or resume the cleaning task.
-
-        This method must be run in the event loop and returns a coroutine.
-        """
-        return self.hass.async_add_job(
-            partial(self.start_pause, **kwargs))
-
-    def send_command(self, command, params=None, **kwargs):
-        """Send a command to a vacuum cleaner."""
-        raise NotImplementedError()
-
-    def async_send_command(self, command, params=None, **kwargs):
-        """Send a command to a vacuum cleaner.
-
-        This method must be run in the event loop and returns a coroutine.
-        """
-        return self.hass.async_add_job(
-            partial(self.send_command, command, params=params, **kwargs))

homeassistant/components/vacuum/demo.py

@@ -10,7 +10,9 @@ from homeassistant.components.vacuum import (
     ATTR_CLEANED_AREA, SUPPORT_BATTERY, SUPPORT_CLEAN_SPOT,
     SUPPORT_FAN_SPEED, SUPPORT_LOCATE, SUPPORT_PAUSE, SUPPORT_RETURN_HOME,
     SUPPORT_SEND_COMMAND, SUPPORT_STATUS, SUPPORT_STOP, SUPPORT_TURN_OFF,
-    SUPPORT_TURN_ON, VacuumDevice)
+    SUPPORT_TURN_ON, SUPPORT_STATE, STATE_CLEANING, STATE_DOCKED,
+    STATE_IDLE, STATE_PAUSED, STATE_RETURNING, VacuumDevice,
+    StateVacuumDevice)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -28,12 +30,17 @@ SUPPORT_ALL_SERVICES = SUPPORT_TURN_ON | SUPPORT_TURN_OFF | SUPPORT_PAUSE | \
                        SUPPORT_LOCATE | SUPPORT_STATUS | SUPPORT_BATTERY | \
                        SUPPORT_CLEAN_SPOT
 
+SUPPORT_STATE_SERVICES = SUPPORT_STATE | SUPPORT_PAUSE | SUPPORT_STOP | \
+                         SUPPORT_RETURN_HOME | SUPPORT_FAN_SPEED | \
+                         SUPPORT_BATTERY | SUPPORT_CLEAN_SPOT
+
 FAN_SPEEDS = ['min', 'medium', 'high', 'max']
 DEMO_VACUUM_COMPLETE = '0_Ground_floor'
 DEMO_VACUUM_MOST = '1_First_floor'
 DEMO_VACUUM_BASIC = '2_Second_floor'
 DEMO_VACUUM_MINIMAL = '3_Third_floor'
 DEMO_VACUUM_NONE = '4_Fourth_floor'
+DEMO_VACUUM_STATE = '5_Fifth_floor'
 
 
 def setup_platform(hass, config, add_devices, discovery_info=None):
@@ -44,6 +51,7 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
         DemoVacuum(DEMO_VACUUM_BASIC, SUPPORT_BASIC_SERVICES),
         DemoVacuum(DEMO_VACUUM_MINIMAL, SUPPORT_MINIMAL_SERVICES),
         DemoVacuum(DEMO_VACUUM_NONE, 0),
+        StateDemoVacuum(DEMO_VACUUM_STATE),
     ])
 
 
@@ -204,3 +212,118 @@ class DemoVacuum(VacuumDevice):
         self._status = 'Executing {}({})'.format(command, params)
         self._state = True
         self.schedule_update_ha_state()
+
+
+class StateDemoVacuum(StateVacuumDevice):
+    """Representation of a demo vacuum supporting states."""
+
+    def __init__(self, name):
+        """Initialize the vacuum."""
+        self._name = name
+        self._supported_features = SUPPORT_STATE_SERVICES
+        self._state = STATE_DOCKED
+        self._fan_speed = FAN_SPEEDS[1]
+        self._cleaned_area = 0
+        self._battery_level = 100
+
+    @property
+    def name(self):
+        """Return the name of the vacuum."""
+        return self._name
+
+    @property
+    def should_poll(self):
+        """No polling needed for a demo vacuum."""
+        return False
+
+    @property
+    def supported_features(self):
+        """Flag supported features."""
+        return self._supported_features
+
+    @property
+    def state(self):
+        """Return the current state of the vacuum."""
+        return self._state
+
+    @property
+    def battery_level(self):
+        """Return the current battery level of the vacuum."""
+        if self.supported_features & SUPPORT_BATTERY == 0:
+            return
+
+        return max(0, min(100, self._battery_level))
+
+    @property
+    def fan_speed(self):
+        """Return the current fan speed of the vacuum."""
+        if self.supported_features & SUPPORT_FAN_SPEED == 0:
+            return
+
+        return self._fan_speed
+
+    @property
+    def fan_speed_list(self):
+        """Return the list of supported fan speeds."""
+        if self.supported_features & SUPPORT_FAN_SPEED == 0:
+            return
+        return FAN_SPEEDS
+
+    @property
+    def device_state_attributes(self):
+        """Return device state attributes."""
+        return {ATTR_CLEANED_AREA: round(self._cleaned_area, 2)}
+
+    def start_pause(self, **kwargs):
+        """Start, pause or resume the cleaning task."""
+        if self.supported_features & SUPPORT_PAUSE == 0:
+            return
+
+        if self._state == STATE_CLEANING:
+            self._state = STATE_PAUSED
+        else:
+            self._state = STATE_CLEANING
+            self._cleaned_area += 1.32
+            self._battery_level -= 1
+        self.schedule_update_ha_state()
+
+    def stop(self, **kwargs):
+        """Stop the cleaning task, do not return to dock."""
+        if self.supported_features & SUPPORT_STOP == 0:
+            return
+
+        self._state = STATE_IDLE
+        self.schedule_update_ha_state()
+
+    def return_to_base(self, **kwargs):
+        """Return dock to charging base."""
+        if self.supported_features & SUPPORT_RETURN_HOME == 0:
+            return
+
+        self._state = STATE_RETURNING
+        self.schedule_update_ha_state()
+
+        self.hass.loop.call_later(30, self.__set_state_to_dock)
+
+    def clean_spot(self, **kwargs):
+        """Perform a spot clean-up."""
+        if self.supported_features & SUPPORT_CLEAN_SPOT == 0:
+            return
+
+        self._state = STATE_CLEANING
+        self._cleaned_area += 1.32
+        self._battery_level -= 1
+        self.schedule_update_ha_state()
+
+    def set_fan_speed(self, fan_speed, **kwargs):
+        """Set the vacuum's fan speed."""
+        if self.supported_features & SUPPORT_FAN_SPEED == 0:
+            return
+
+        if fan_speed in self.fan_speed_list:
+            self._fan_speed = fan_speed
+            self.schedule_update_ha_state()
+
+    def __set_state_to_dock(self):
+        self._state = STATE_DOCKED
+        self.schedule_update_ha_state()

tests/components/vacuum/test_demo.py

@@ -6,10 +6,12 @@ from homeassistant.components.vacuum import (
     ATTR_BATTERY_LEVEL, ATTR_COMMAND, ATTR_ENTITY_ID, ATTR_FAN_SPEED,
     ATTR_FAN_SPEED_LIST, ATTR_PARAMS, ATTR_STATUS, DOMAIN,
     ENTITY_ID_ALL_VACUUMS,
-    SERVICE_SEND_COMMAND, SERVICE_SET_FAN_SPEED)
+    SERVICE_SEND_COMMAND, SERVICE_SET_FAN_SPEED,
+    STATE_DOCKED, STATE_CLEANING, STATE_PAUSED, STATE_IDLE,
+    STATE_RETURNING)
 from homeassistant.components.vacuum.demo import (
     DEMO_VACUUM_BASIC, DEMO_VACUUM_COMPLETE, DEMO_VACUUM_MINIMAL,
-    DEMO_VACUUM_MOST, DEMO_VACUUM_NONE, FAN_SPEEDS)
+    DEMO_VACUUM_MOST, DEMO_VACUUM_NONE, DEMO_VACUUM_STATE, FAN_SPEEDS)
 from homeassistant.const import (
     ATTR_SUPPORTED_FEATURES, CONF_PLATFORM, STATE_OFF, STATE_ON)
 from homeassistant.setup import setup_component
@@ -21,6 +23,7 @@ ENTITY_VACUUM_COMPLETE = '{}.{}'.format(DOMAIN, DEMO_VACUUM_COMPLETE).lower()
 ENTITY_VACUUM_MINIMAL = '{}.{}'.format(DOMAIN, DEMO_VACUUM_MINIMAL).lower()
 ENTITY_VACUUM_MOST = '{}.{}'.format(DOMAIN, DEMO_VACUUM_MOST).lower()
 ENTITY_VACUUM_NONE = '{}.{}'.format(DOMAIN, DEMO_VACUUM_NONE).lower()
+ENTITY_VACUUM_STATE = '{}.{}'.format(DOMAIN, DEMO_VACUUM_STATE).lower()
 
 
 class TestVacuumDemo(unittest.TestCase):
@@ -79,6 +82,14 @@ class TestVacuumDemo(unittest.TestCase):
         self.assertEqual(None, state.attributes.get(ATTR_FAN_SPEED_LIST))
         self.assertEqual(STATE_OFF, state.state)
 
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertEqual(5244, state.attributes.get(ATTR_SUPPORTED_FEATURES))
+        self.assertEqual(STATE_DOCKED, state.state)
+        self.assertEqual(100, state.attributes.get(ATTR_BATTERY_LEVEL))
+        self.assertEqual("medium", state.attributes.get(ATTR_FAN_SPEED))
+        self.assertListEqual(FAN_SPEEDS,
+                             state.attributes.get(ATTR_FAN_SPEED_LIST))
+
     def test_methods(self):
         """Test if methods call the services as expected."""
         self.hass.states.set(ENTITY_VACUUM_BASIC, STATE_ON)
@@ -147,6 +158,41 @@ class TestVacuumDemo(unittest.TestCase):
         self.assertIn("spot", state.attributes.get(ATTR_STATUS))
         self.assertEqual(STATE_ON, state.state)
 
+        vacuum.start_pause(self.hass, ENTITY_VACUUM_STATE)
+        self.hass.block_till_done()
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertEqual(STATE_CLEANING, state.state)
+
+        vacuum.start_pause(self.hass, ENTITY_VACUUM_STATE)
+        self.hass.block_till_done()
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertEqual(STATE_PAUSED, state.state)
+
+        vacuum.stop(self.hass, ENTITY_VACUUM_STATE)
+        self.hass.block_till_done()
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertEqual(STATE_IDLE, state.state)
+
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertLess(state.attributes.get(ATTR_BATTERY_LEVEL), 100)
+        self.assertNotEqual(STATE_DOCKED, state.state)
+
+        vacuum.return_to_base(self.hass, ENTITY_VACUUM_STATE)
+        self.hass.block_till_done()
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertEqual(STATE_RETURNING, state.state)
+
+        vacuum.set_fan_speed(self.hass, FAN_SPEEDS[-1],
+                             entity_id=ENTITY_VACUUM_STATE)
+        self.hass.block_till_done()
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertEqual(FAN_SPEEDS[-1], state.attributes.get(ATTR_FAN_SPEED))
+
+        vacuum.clean_spot(self.hass, entity_id=ENTITY_VACUUM_STATE)
+        self.hass.block_till_done()
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertEqual(STATE_CLEANING, state.state)
+
     def test_unsupported_methods(self):
         """Test service calls for unsupported vacuums."""
         self.hass.states.set(ENTITY_VACUUM_NONE, STATE_ON)
@@ -201,6 +247,22 @@ class TestVacuumDemo(unittest.TestCase):
         self.assertNotIn("spot", state.attributes.get(ATTR_STATUS))
         self.assertEqual(STATE_OFF, state.state)
 
+        # StateVacuumDevice does not support on/off
+        vacuum.turn_on(self.hass, entity_id=ENTITY_VACUUM_STATE)
+        self.hass.block_till_done()
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertNotEqual(STATE_CLEANING, state.state)
+
+        vacuum.turn_off(self.hass, entity_id=ENTITY_VACUUM_STATE)
+        self.hass.block_till_done()
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertNotEqual(STATE_RETURNING, state.state)
+
+        vacuum.toggle(self.hass, entity_id=ENTITY_VACUUM_STATE)
+        self.hass.block_till_done()
+        state = self.hass.states.get(ENTITY_VACUUM_STATE)
+        self.assertNotEqual(STATE_CLEANING, state.state)
+
     def test_services(self):
         """Test vacuum services."""
         # Test send_command
@@ -241,9 +303,11 @@ class TestVacuumDemo(unittest.TestCase):
     def test_set_fan_speed(self):
         """Test vacuum service to set the fan speed."""
         group_vacuums = ','.join([ENTITY_VACUUM_BASIC,
-                                  ENTITY_VACUUM_COMPLETE])
+                                  ENTITY_VACUUM_COMPLETE,
+                                  ENTITY_VACUUM_STATE])
         old_state_basic = self.hass.states.get(ENTITY_VACUUM_BASIC)
         old_state_complete = self.hass.states.get(ENTITY_VACUUM_COMPLETE)
+        old_state_state = self.hass.states.get(ENTITY_VACUUM_STATE)
 
         vacuum.set_fan_speed(
             self.hass, FAN_SPEEDS[0], entity_id=group_vacuums)
@@ -251,6 +315,7 @@ class TestVacuumDemo(unittest.TestCase):
         self.hass.block_till_done()
         new_state_basic = self.hass.states.get(ENTITY_VACUUM_BASIC)
         new_state_complete = self.hass.states.get(ENTITY_VACUUM_COMPLETE)
+        new_state_state = self.hass.states.get(ENTITY_VACUUM_STATE)
 
         self.assertEqual(old_state_basic, new_state_basic)
         self.assertNotIn(ATTR_FAN_SPEED, new_state_basic.attributes)
@@ -261,6 +326,12 @@ class TestVacuumDemo(unittest.TestCase):
         self.assertEqual(FAN_SPEEDS[0],
                          new_state_complete.attributes[ATTR_FAN_SPEED])
 
+        self.assertNotEqual(old_state_state, new_state_state)
+        self.assertEqual(FAN_SPEEDS[1],
+                         old_state_state.attributes[ATTR_FAN_SPEED])
+        self.assertEqual(FAN_SPEEDS[0],
+                         new_state_state.attributes[ATTR_FAN_SPEED])
+
     def test_send_command(self):
         """Test vacuum service to send a command."""
         group_vacuums = ','.join([ENTITY_VACUUM_BASIC,