Implement external state restoration (#58)

* Consolidate init-time state reading for Light entity * Read the color mode from cache as well * Remove duplicate `supported_color_modes` variable * Switch `zcl_color_mode_to_entity_color_mode` to a static dictionary * Only set the color mode from the supported color modes if it is uncached * Update the attribute cache color mode after the color has been successfully set * Do not persist the color mode for groups * Test that the color mode changes * Account for invalid ZCL color modes * Add a quick test for HS * Unit test enhanced hue as well * Re-introduce erroneously removed `cached_property` * Add `restore_extra_state_attributes` * Persist the door lock state after locking/unlocking * Remove unused lock `kwargs` * Add `restore_external_state_attributes` * Implement external state for `cover` * Implement external state for `select` * Implement external state for `siren` * Remove unnecessary `_persist_lock_state` * Revert "Implement external state for `siren`" This reverts commit 7ef8ba3. * Implement a stub `restore_external_state_attributes` for non-ZCL selects * Migrate coverage to `pyproject.toml` and exclude NotImplementedError * Update zha/application/platforms/light/__init__.py Co-authored-by: TheJulianJES <[email protected]> * Migrate lighting to use explicit state restoration instead of ZCL cache * Reduce diff size * Only restore the state if the attribute isn't `None` * Migrate lock to use state restoration * Add some unit tests * Offload validation to Core * Implement an `undefined` type * Migrate remaining platforms to use `UNDEFINED` as well, where appropriate * Finish unit tests * Remove `UNDEFINED` * Only restore (most) light state attributes if they are not `None` * Fix `number` entity name * Revert `cached_property` -> `property` change --------- Co-authored-by: TheJulianJES <[email protected]>
zigpy · Jul 5, 2024 · 8abf678 · 8abf678
1 parent 4fbc162
commit 8abf678
Show file tree

Hide file tree

Showing 15 changed files with 290 additions and 23 deletions.
diff --git a/.coveragerc b/.coveragerc
diff --git a/pyproject.toml b/pyproject.toml
@@ -231,4 +231,11 @@ split-on-trailing-comma = false
 "script/*" = ["T20"]
 
 [tool.ruff.lint.mccabe]
-max-complexity = 25
+max-complexity = 25
+
+[tool.coverage.report]
+show_missing = true
+exclude_also = [
+    "if TYPE_CHECKING:",
+    "raise NotImplementedError",
+]
diff --git a/requirements_test.txt b/requirements_test.txt
@@ -1,4 +1,4 @@
-codecov
+coverage[toml]
 colorlog
 codespell
 mypy

diff --git a/tests/test_cover.py b/tests/test_cover.py
@@ -799,3 +799,27 @@ async def test_cover_remote(
     assert zha_device.emit_zha_event.call_count == 1
     assert ATTR_COMMAND in zha_device.emit_zha_event.call_args[0][0]
     assert zha_device.emit_zha_event.call_args[0][0][ATTR_COMMAND] == "down_close"
+
+
+async def test_cover_state_restoration(
+    device_joined: Callable[[ZigpyDevice], Awaitable[Device]],
+    zigpy_cover_device: ZigpyDevice,
+    zha_gateway: Gateway,
+) -> None:
+    """Test the cover state restoration."""
+    zha_device = await device_joined(zigpy_cover_device)
+    entity = get_entity(zha_device, platform=Platform.COVER)
+
+    assert entity.state["state"] != STATE_CLOSED
+    assert entity.state["target_lift_position"] != 12
+    assert entity.state["target_tilt_position"] != 34
+
+    entity.restore_external_state_attributes(
+        state=STATE_CLOSED,
+        target_lift_position=12,
+        target_tilt_position=34,
+    )
+
+    assert entity.state["state"] == STATE_CLOSED
+    assert entity.state["target_lift_position"] == 12
+    assert entity.state["target_tilt_position"] == 34
diff --git a/tests/test_light.py b/tests/test_light.py
@@ -209,6 +209,15 @@ async def device_light_3(
         ieee=IEEE_GROUPABLE_DEVICE3,
         nwk=0xB89F,
     )
+    color_cluster = zigpy_device.endpoints[1].light_color
+    color_cluster.PLUGGED_ATTR_READS = {
+        "color_capabilities": (
+            lighting.Color.ColorCapabilities.Color_temperature
+            | lighting.Color.ColorCapabilities.XY_attributes
+            | lighting.Color.ColorCapabilities.Color_loop
+        )
+    }
+
     zha_device = await device_joined(zigpy_device)
     zha_device.available = True
     return zha_device
@@ -242,8 +251,10 @@ async def eWeLink_light(
     )
     color_cluster = zigpy_device.endpoints[1].light_color
     color_cluster.PLUGGED_ATTR_READS = {
-        "color_capabilities": lighting.Color.ColorCapabilities.Color_temperature
-        | lighting.Color.ColorCapabilities.XY_attributes,
+        "color_capabilities": (
+            lighting.Color.ColorCapabilities.Color_temperature
+            | lighting.Color.ColorCapabilities.XY_attributes
+        ),
         "color_temp_physical_min": 0,
         "color_temp_physical_max": 0,
     }
@@ -334,8 +345,11 @@ async def test_light(
             "color_temperature": 100,
             "color_temp_physical_min": 0,
             "color_temp_physical_max": 600,
-            "color_capabilities": lighting.ColorCapabilities.XY_attributes
-            | lighting.ColorCapabilities.Color_temperature,
+            "color_capabilities": (
+                lighting.ColorCapabilities.XY_attributes
+                | lighting.ColorCapabilities.Color_temperature
+                | lighting.ColorCapabilities.Hue_and_saturation
+            ),
         }
         update_attribute_cache(cluster_color)
     zha_device = await device_joined(zigpy_device)
@@ -398,6 +412,7 @@ async def test_light(
         assert entity.state["color_temp"] != 200
         await entity.async_turn_on(brightness=50, transition=10, color_temp=200)
         await zha_gateway.async_block_till_done()
+        assert entity.state["color_mode"] == ColorMode.COLOR_TEMP
         assert entity.state["brightness"] == 50
         assert entity.state["color_temp"] == 200
         assert bool(entity.state["on"]) is True
@@ -419,6 +434,7 @@ async def test_light(
         assert entity.state["xy_color"] != [13369, 18087]
         await entity.async_turn_on(brightness=50, xy_color=[13369, 18087])
         await zha_gateway.async_block_till_done()
+        assert entity.state["color_mode"] == ColorMode.XY
         assert entity.state["brightness"] == 50
         assert entity.state["xy_color"] == [13369, 18087]
         assert cluster_color.request.call_count == 1
@@ -437,6 +453,58 @@ async def test_light(
 
         cluster_color.request.reset_mock()
 
+        # test color hs from the client
+        assert entity.state["hs_color"] != [12, 34]
+        await entity.async_turn_on(brightness=50, hs_color=[12, 34])
+        await zha_gateway.async_block_till_done()
+        assert entity.state["color_mode"] == ColorMode.HS
+        assert entity.state["brightness"] == 50
+        assert entity.state["hs_color"] == [12, 34]
+        assert cluster_color.request.call_count == 1
+        assert cluster_color.request.await_count == 1
+        assert cluster_color.request.call_args == call(
+            False,
+            6,
+            cluster_color.commands_by_name["move_to_hue_and_saturation"].schema,
+            hue=8,
+            saturation=86,
+            transition_time=0,
+            expect_reply=True,
+            manufacturer=None,
+            tsn=None,
+        )
+
+        cluster_color.request.reset_mock()
+
+        # test enhanced hue support
+        cluster_color.PLUGGED_ATTR_READS["color_capabilities"] |= (
+            lighting.ColorCapabilities.Enhanced_hue
+        )
+        update_attribute_cache(cluster_color)
+        del entity._color_cluster_handler.color_capabilities
+
+        assert entity.state["hs_color"] != [56, 78]
+        await entity.async_turn_on(brightness=50, hs_color=[56, 78])
+        await zha_gateway.async_block_till_done()
+        assert entity.state["color_mode"] == ColorMode.HS
+        assert entity.state["brightness"] == 50
+        assert entity.state["hs_color"] == [56, 78]
+        assert cluster_color.request.call_count == 1
+        assert cluster_color.request.await_count == 1
+        assert cluster_color.request.call_args == call(
+            False,
+            67,
+            cluster_color.commands_by_name[
+                "enhanced_move_to_hue_and_saturation"
+            ].schema,
+            enhanced_hue=10194,
+            saturation=198,
+            transition_time=0,
+            expect_reply=True,
+            manufacturer=None,
+            tsn=None,
+        )
+
 
 async def async_test_on_off_from_light(
     zha_gateway: Gateway,
@@ -1836,3 +1904,47 @@ async def test_group_member_assume_state(
     assert bool(device_1_light_entity.state["on"]) is False
     assert bool(device_2_light_entity.state["on"]) is False
     assert bool(entity.state["on"]) is False
+
+
+async def test_light_state_restoration(
+    device_light_3,  # pylint: disable=redefined-outer-name
+) -> None:
+    """Test the light state restoration function."""
+    entity = get_entity(device_light_3, platform=Platform.LIGHT)
+    entity.restore_external_state_attributes(
+        state=True,
+        off_with_transition=False,
+        off_brightness=12,
+        brightness=34,
+        color_temp=500,
+        xy_color=(1, 2),
+        hs_color=(3, 4),
+        color_mode=ColorMode.XY,
+        effect="colorloop",
+    )
+
+    assert entity.state["on"] is True
+    assert entity.state["brightness"] == 34
+    assert entity.state["color_temp"] == 500
+    assert entity.state["xy_color"] == (1, 2)
+    assert entity.state["color_mode"] == ColorMode.XY
+    assert entity.state["effect"] == "colorloop"
+
+    entity.restore_external_state_attributes(
+        state=None,
+        off_with_transition=None,
+        off_brightness=None,
+        brightness=None,
+        color_temp=None,
+        xy_color=None,
+        hs_color=None,
+        color_mode=None,
+        effect=None,  # Effect is the only `None` value actually restored
+    )
+
+    assert entity.state["on"] is True
+    assert entity.state["brightness"] == 34
+    assert entity.state["color_temp"] == 500
+    assert entity.state["xy_color"] == (1, 2)
+    assert entity.state["color_mode"] == ColorMode.XY
+    assert entity.state["effect"] is None
diff --git a/tests/test_lock.py b/tests/test_lock.py
@@ -14,6 +14,7 @@
 from zha.application import Platform
 from zha.application.gateway import Gateway
 from zha.application.platforms import PlatformEntity
+from zha.application.platforms.lock.const import STATE_LOCKED, STATE_UNLOCKED
 from zha.zigbee.device import Device
 
 LOCK_DOOR = 0
@@ -210,3 +211,20 @@ async def async_disable_user_code(
         assert cluster.request.call_args[0][1] == SET_USER_STATUS
         assert cluster.request.call_args[0][3] == 2  # user slot 3 => internal slot 2
         assert cluster.request.call_args[0][4] == closures.DoorLock.UserStatus.Disabled
+
+
+async def test_lock_state_restoration(
+    lock: tuple[Device, closures.DoorLock],  # pylint: disable=redefined-outer-name
+    zha_gateway: Gateway,
+) -> None:
+    """Test the lock state restoration."""
+    zha_device, _ = lock
+    entity = get_entity(zha_device, platform=Platform.LOCK)
+
+    assert entity.state["is_locked"] is False
+
+    entity.restore_external_state_attributes(state=STATE_LOCKED)
+    assert entity.state["is_locked"] is True
+
+    entity.restore_external_state_attributes(state=STATE_UNLOCKED)
+    assert entity.state["is_locked"] is False
diff --git a/tests/test_number.py b/tests/test_number.py
@@ -102,6 +102,8 @@ async def test_number(
 
     assert cluster.read_attributes.call_count == 3
 
+    assert entity.name == "PWM1"
+
     # test that the state is 15.0
     assert entity.state["state"] == 15.0
 

diff --git a/tests/test_select.py b/tests/test_select.py
@@ -238,3 +238,24 @@ async def test_on_off_select_attribute_report_v2(
     assert cluster.write_attributes.call_args == call(
         {"motion_sensitivity": AqaraMotionSensitivities.Medium}, manufacturer=None
     )
+
+
+async def test_non_zcl_select_state_restoration(
+    siren: tuple[Device, security.IasWd],  # pylint: disable=redefined-outer-name
+    zha_gateway: Gateway,
+) -> None:
+    """Test the non-ZCL select state restoration."""
+    zha_device, cluster = siren
+    entity = get_entity(zha_device, platform=Platform.SELECT, qualifier="WarningMode")
+
+    assert entity.state["state"] is None
+
+    entity.restore_external_state_attributes(
+        state=security.IasWd.Warning.WarningMode.Burglar.name
+    )
+    assert entity.state["state"] == security.IasWd.Warning.WarningMode.Burglar.name
+
+    entity.restore_external_state_attributes(
+        state=security.IasWd.Warning.WarningMode.Fire.name
+    )
+    assert entity.state["state"] == security.IasWd.Warning.WarningMode.Fire.name
diff --git a/zha/application/platforms/__init__.py b/zha/application/platforms/__init__.py
@@ -224,6 +224,13 @@ def extra_state_attribute_names(self) -> set[str] | None:
             return self._attr_extra_state_attribute_names
         return None
 
+    def restore_external_state_attributes(self, **kwargs: Any) -> None:
+        """Restore entity specific state attributes from an external source.
+
+        Entities implementing this must accept a keyword argument for each attribute.
+        """
+        raise NotImplementedError
+
     async def on_remove(self) -> None:
         """Cancel tasks and timers this entity owns."""
         for handle in self._tracked_handles:

diff --git a/zha/application/platforms/cover/__init__.py b/zha/application/platforms/cover/__init__.py
@@ -5,7 +5,7 @@
 import asyncio
 import functools
 import logging
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from zigpy.zcl.clusters.general import OnOff
 from zigpy.zcl.foundation import Status
@@ -57,6 +57,10 @@ class Cover(PlatformEntity):
     PLATFORM = Platform.COVER
 
     _attr_translation_key: str = "cover"
+    _attr_extra_state_attribute_names: set[str] = {
+        "target_lift_position",
+        "target_tilt_position",
+    }
 
     def __init__(
         self,
@@ -84,7 +88,7 @@ def __init__(
         )
         self._target_lift_position: int | None = None
         self._target_tilt_position: int | None = None
-        self._state: str
+        self._state: str = STATE_OPEN
         self._determine_initial_state()
         self._cover_cluster_handler.on_event(
             CLUSTER_HANDLER_ATTRIBUTE_UPDATED,
@@ -102,10 +106,26 @@ def state(self) -> dict[str, Any]:
                 "is_opening": self.is_opening,
                 "is_closing": self.is_closing,
                 "is_closed": self.is_closed,
+                "target_lift_position": self._target_lift_position,
+                "target_tilt_position": self._target_tilt_position,
             }
         )
         return response
 
+    def restore_external_state_attributes(
+        self,
+        *,
+        state: Literal[
+            "open", "opening", "closed", "closing"
+        ],  # FIXME: why must these be expanded?
+        target_lift_position: int | None,
+        target_tilt_position: int | None,
+    ):
+        """Restore external state attributes."""
+        self._state = state
+        self._target_lift_position = target_lift_position
+        self._target_tilt_position = target_tilt_position
+
     @property
     def is_closed(self) -> bool | None:
         """Return True if the cover is closed.

diff --git a/zha/application/platforms/cover/const.py b/zha/application/platforms/cover/const.py
@@ -10,10 +10,10 @@
 ATTR_POSITION: Final[str] = "position"
 ATTR_TILT_POSITION: Final[str] = "tilt_position"
 
-STATE_OPEN: Final[str] = "open"
-STATE_OPENING: Final[str] = "opening"
-STATE_CLOSED: Final[str] = "closed"
-STATE_CLOSING: Final[str] = "closing"
+STATE_OPEN: Final = "open"
+STATE_OPENING: Final = "opening"
+STATE_CLOSED: Final = "closed"
+STATE_CLOSING: Final = "closing"
 
 
 class CoverDeviceClass(StrEnum):