Do optimistic state update for Z-Wave multilevel switch entities (#90490)

* Do optimistic state update for Z-Wave multilevel switch entities

* simplify

* define constant for setting value to previous value

* Rework to only consider value of 255 and only places where we know that the previous state is known by the integration due to the service being called

* missed commit

* better code

* Add tests and use constant from lib

* fix logic

* fix bug

* Add comments with more details

homeassistant/components/zwave_js/fan.py

@@ -6,6 +6,7 @@ from typing import Any, cast
 
 from zwave_js_server.client import Client as ZwaveClient
 from zwave_js_server.const import TARGET_VALUE_PROPERTY, CommandClass
+from zwave_js_server.const.command_class.multilevel_switch import SET_TO_PREVIOUS_VALUE
 from zwave_js_server.const.command_class.thermostat import (
     THERMOSTAT_FAN_OFF_PROPERTY,
     THERMOSTAT_FAN_STATE_PROPERTY,
@@ -88,6 +89,8 @@ class ZwaveFan(ZWaveBaseEntity, FanEntity):
         assert target_value
         self._target_value = target_value
 
+        self._use_optimistic_state: bool = False
+
     async def async_set_percentage(self, percentage: int) -> None:
         """Set the speed percentage of the fan."""
         if percentage == 0:
@@ -111,8 +114,19 @@ class ZwaveFan(ZWaveBaseEntity, FanEntity):
         elif preset_mode is not None:
             await self.async_set_preset_mode(preset_mode)
         else:
-            # Value 255 tells device to return to previous value
-            await self.info.node.async_set_value(self._target_value, 255)
+            if self.info.primary_value.command_class != CommandClass.SWITCH_MULTILEVEL:
+                raise HomeAssistantError(
+                    "`percentage` or `preset_mode` must be provided"
+                )
+            # If this is a Multilevel Switch CC value, we do an optimistic state update
+            # when setting to a previous value to avoid waiting for the value to be
+            # updated from the device which is typically delayed and causes a confusing
+            # UX.
+            await self.info.node.async_set_value(
+                self._target_value, SET_TO_PREVIOUS_VALUE
+            )
+            self._use_optimistic_state = True
+            self.async_write_ha_state()
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the device off."""
@@ -121,6 +135,9 @@ class ZwaveFan(ZWaveBaseEntity, FanEntity):
     @property
     def is_on(self) -> bool | None:
         """Return true if device is on (speed above 0)."""
+        if self._use_optimistic_state:
+            self._use_optimistic_state = False
+            return True
         if self.info.primary_value.value is None:
             # guard missing value
             return None

homeassistant/components/zwave_js/light.py

@@ -22,6 +22,7 @@ from zwave_js_server.const.command_class.color_switch import (
     TARGET_COLOR_PROPERTY,
     ColorComponent,
 )
+from zwave_js_server.const.command_class.multilevel_switch import SET_TO_PREVIOUS_VALUE
 from zwave_js_server.model.driver import Driver
 from zwave_js_server.model.value import Value
 
@@ -164,6 +165,8 @@ class ZwaveLight(ZWaveBaseEntity, LightEntity):
         if self.supports_brightness_transition or self.supports_color_transition:
             self._attr_supported_features |= LightEntityFeature.TRANSITION
 
+        self._set_optimistic_state: bool = False
+
     @callback
     def on_value_update(self) -> None:
         """Call when a watched value is added or updated."""
@@ -187,10 +190,11 @@ class ZwaveLight(ZWaveBaseEntity, LightEntity):
     @property
     def is_on(self) -> bool | None:
         """Return true if device is on (brightness above 0)."""
+        if self._set_optimistic_state:
+            self._set_optimistic_state = False
+            return True
         brightness = self.brightness
-        if brightness is None:
-            return None
-        return brightness > 0
+        return brightness > 0 if brightness is not None else None
 
     @property
     def hs_color(self) -> tuple[float, float] | None:
@@ -332,8 +336,7 @@ class ZwaveLight(ZWaveBaseEntity, LightEntity):
         if not self._target_brightness:
             return
         if brightness is None:
-            # Level 255 means to set it to previous value.
-            zwave_brightness = 255
+            zwave_brightness = SET_TO_PREVIOUS_VALUE
         else:
             # Zwave multilevel switches use a range of [0, 99] to control brightness.
             zwave_brightness = byte_to_zwave_brightness(brightness)
@@ -350,6 +353,15 @@ class ZwaveLight(ZWaveBaseEntity, LightEntity):
         await self.info.node.async_set_value(
             self._target_brightness, zwave_brightness, zwave_transition
         )
+        # We do an optimistic state update when setting to a previous value
+        # to avoid waiting for the value to be updated from the device which is
+        # typically delayed and causes a confusing UX.
+        if (
+            zwave_brightness == SET_TO_PREVIOUS_VALUE
+            and self.info.primary_value.command_class == CommandClass.SWITCH_MULTILEVEL
+        ):
+            self._set_optimistic_state = True
+            self.async_write_ha_state()
 
     @callback
     def _calculate_color_values(self) -> None:

tests/components/zwave_js/test_fan.py

@@ -43,7 +43,35 @@ async def test_generic_fan(
     state = hass.states.get(entity_id)
 
     assert state
-    assert state.state == "off"
+    assert state.state == STATE_OFF
+
+    # Test turn on no speed
+    await hass.services.async_call(
+        "fan",
+        "turn_on",
+        {"entity_id": entity_id},
+        blocking=True,
+    )
+
+    assert len(client.async_send_command.call_args_list) == 1
+    args = client.async_send_command.call_args[0][0]
+    assert args["command"] == "node.set_value"
+    assert args["nodeId"] == 17
+    assert args["valueId"] == {
+        "commandClass": 38,
+        "endpoint": 0,
+        "property": "targetValue",
+    }
+    assert args["value"] == 255
+
+    client.async_send_command.reset_mock()
+
+    # Due to optimistic updates, the state should be on even though the Z-Wave state
+    # hasn't been updated yet
+    state = hass.states.get(entity_id)
+
+    assert state
+    assert state.state == STATE_ON
 
     # Test turn on setting speed
     await hass.services.async_call(
@@ -77,27 +105,6 @@ async def test_generic_fan(
 
     client.async_send_command.reset_mock()
 
-    # Test turn on no speed
-    await hass.services.async_call(
-        "fan",
-        "turn_on",
-        {"entity_id": entity_id},
-        blocking=True,
-    )
-
-    assert len(client.async_send_command.call_args_list) == 1
-    args = client.async_send_command.call_args[0][0]
-    assert args["command"] == "node.set_value"
-    assert args["nodeId"] == 17
-    assert args["valueId"] == {
-        "commandClass": 38,
-        "endpoint": 0,
-        "property": "targetValue",
-    }
-    assert args["value"] == 255
-
-    client.async_send_command.reset_mock()
-
     # Test turning off
     await hass.services.async_call(
         "fan",
@@ -140,7 +147,7 @@ async def test_generic_fan(
     node.receive_event(event)
 
     state = hass.states.get(entity_id)
-    assert state.state == "on"
+    assert state.state == STATE_ON
     assert state.attributes[ATTR_PERCENTAGE] == 100
 
     client.async_send_command.reset_mock()
@@ -165,7 +172,7 @@ async def test_generic_fan(
     node.receive_event(event)
 
     state = hass.states.get(entity_id)
-    assert state.state == "off"
+    assert state.state == STATE_OFF
     assert state.attributes[ATTR_PERCENTAGE] == 0
 
     client.async_send_command.reset_mock()

tests/components/zwave_js/test_light.py

@@ -70,6 +70,13 @@ async def test_light(
     }
     assert args["value"] == 255
 
+    # Due to optimistic updates, the state should be on even though the Z-Wave state
+    # hasn't been updated yet
+    state = hass.states.get(BULB_6_MULTI_COLOR_LIGHT_ENTITY)
+
+    assert state
+    assert state.state == STATE_ON
+
     client.async_send_command.reset_mock()
 
     # Test turning on with transition