From a58b3721edb09cd1b724fc241b539bbd63a6ebe7 Mon Sep 17 00:00:00 2001
From: TheJulianJES <TheJulianJES@users.noreply.github.com>
Date: Tue, 4 Apr 2023 04:27:57 +0200
Subject: [PATCH] Restore state for ZHA OnOff binary sensors (#90749)

* Restore state for ZHA OnOff binary sensors

* Let `Motion` extend `Opening`

`Motion` is just a specified version of `Opening` that only changes the device class for some motion sensors.
Since we have more "special code" in the OnOff/Opening sensor now, we also want to make sure that gets applied to `Motion` binary sensors.

* Improve comment and type

* Add test to verify that binary sensors restore last HA state
---
 homeassistant/components/zha/binary_sensor.py | 19 ++++++++--
 tests/components/zha/test_binary_sensor.py    | 38 +++++++++++++++++++
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/zha/binary_sensor.py b/homeassistant/components/zha/binary_sensor.py
index 4e3c7166bf0..696216e3e81 100644
--- a/homeassistant/components/zha/binary_sensor.py
+++ b/homeassistant/components/zha/binary_sensor.py
@@ -4,6 +4,8 @@ from __future__ import annotations
 import functools
 from typing import Any
 
+import zigpy.types as t
+from zigpy.zcl.clusters.general import OnOff
 from zigpy.zcl.clusters.security import IasZone
 
 from homeassistant.components.binary_sensor import (
@@ -119,11 +121,21 @@ class Occupancy(BinarySensor):
 
 @STRICT_MATCH(channel_names=CHANNEL_ON_OFF)
 class Opening(BinarySensor):
-    """ZHA BinarySensor."""
+    """ZHA OnOff BinarySensor."""
 
     SENSOR_ATTR = "on_off"
     _attr_device_class: BinarySensorDeviceClass = BinarySensorDeviceClass.OPENING
 
+    # Client/out cluster attributes aren't stored in the zigpy database, but are properly stored in the runtime cache.
+    # We need to manually restore the last state from the sensor state to the runtime cache for now.
+    @callback
+    def async_restore_last_state(self, last_state):
+        """Restore previous state to zigpy cache."""
+        self._channel.cluster.update_attribute(
+            OnOff.attributes_by_name[self.SENSOR_ATTR].id,
+            t.Bool.true if last_state.state == STATE_ON else t.Bool.false,
+        )
+
 
 @MULTI_MATCH(channel_names=CHANNEL_BINARY_INPUT)
 class BinaryInput(BinarySensor):
@@ -144,10 +156,9 @@ class BinaryInput(BinarySensor):
     manufacturers="Philips",
     models={"SML001", "SML002"},
 )
-class Motion(BinarySensor):
-    """ZHA BinarySensor."""
+class Motion(Opening):
+    """ZHA OnOff BinarySensor with motion device class."""
 
-    SENSOR_ATTR = "on_off"
     _attr_device_class: BinarySensorDeviceClass = BinarySensorDeviceClass.MOTION
 
 
diff --git a/tests/components/zha/test_binary_sensor.py b/tests/components/zha/test_binary_sensor.py
index ec25295ed5a..2c0461a3c7c 100644
--- a/tests/components/zha/test_binary_sensor.py
+++ b/tests/components/zha/test_binary_sensor.py
@@ -3,6 +3,7 @@ from unittest.mock import patch
 
 import pytest
 import zigpy.profiles.zha
+import zigpy.zcl.clusters.general as general
 import zigpy.zcl.clusters.measurement as measurement
 import zigpy.zcl.clusters.security as security
 
@@ -40,6 +41,16 @@ DEVICE_OCCUPANCY = {
 }
 
 
+DEVICE_ONOFF = {
+    1: {
+        SIG_EP_PROFILE: zigpy.profiles.zha.PROFILE_ID,
+        SIG_EP_TYPE: zigpy.profiles.zha.DeviceType.ON_OFF_SENSOR,
+        SIG_EP_INPUT: [],
+        SIG_EP_OUTPUT: [general.OnOff.cluster_id],
+    }
+}
+
+
 @pytest.fixture(autouse=True)
 def binary_sensor_platform_only():
     """Only set up the binary_sensor and required base platforms to speed up tests."""
@@ -212,3 +223,30 @@ async def test_binary_sensor_migration_already_migrated(
     assert entity_id is not None
     assert hass.states.get(entity_id).state == STATE_ON  # matches attribute cache
     assert hass.states.get(entity_id).attributes["migrated_to_cache"]
+
+
+@pytest.mark.parametrize(
+    "restored_state",
+    [
+        STATE_ON,
+        STATE_OFF,
+    ],
+)
+async def test_onoff_binary_sensor_restore_state(
+    hass: HomeAssistant,
+    zigpy_device_mock,
+    core_rs,
+    zha_device_restored,
+    restored_state,
+) -> None:
+    """Test ZHA OnOff binary_sensor restores last state from HA."""
+
+    entity_id = "binary_sensor.fakemanufacturer_fakemodel_opening"
+    core_rs(entity_id, state=restored_state, attributes={})
+
+    zigpy_device = zigpy_device_mock(DEVICE_ONOFF)
+    zha_device = await zha_device_restored(zigpy_device)
+    entity_id = await find_entity_id(Platform.BINARY_SENSOR, zha_device, hass)
+
+    assert entity_id is not None
+    assert hass.states.get(entity_id).state == restored_state