Migrate old ZHA IasZone sensor state to zigpy cache (#90508)

* Migrate old ZHA IasZone sensor state to zigpy cache

* Use correct type for ZoneStatus

* Test that migration happens

* Test that migration only happens once

* Fix parametrize

homeassistant/components/zha/binary_sensor.py

@@ -2,13 +2,16 @@
 from __future__ import annotations
 
 import functools
+from typing import Any
+
+from zigpy.zcl.clusters.security import IasZone
 
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import EntityCategory, Platform
+from homeassistant.const import STATE_ON, EntityCategory, Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -164,6 +167,36 @@ class IASZone(BinarySensor):
         """Parse the raw attribute into a bool state."""
         return BinarySensor.parse(value & 3)  # use only bit 0 and 1 for alarm state
 
+    # temporary code to migrate old IasZone sensors to update attribute cache state once
+    # remove in 2024.4.0
+    @property
+    def extra_state_attributes(self) -> dict[str, Any]:
+        """Return state attributes."""
+        return {"migrated_to_cache": True}  # writing new state means we're migrated
+
+    # temporary migration code
+    @callback
+    def async_restore_last_state(self, last_state):
+        """Restore previous state."""
+        # trigger migration if extra state attribute is not present
+        if "migrated_to_cache" not in last_state.attributes:
+            self.migrate_to_zigpy_cache(last_state)
+
+    # temporary migration code
+    @callback
+    def migrate_to_zigpy_cache(self, last_state):
+        """Save old IasZone sensor state to attribute cache."""
+        # previous HA versions did not update the attribute cache for IasZone sensors, so do it once here
+        # a HA state write is triggered shortly afterwards and writes the "migrated_to_cache" extra state attribute
+        if last_state.state == STATE_ON:
+            migrated_state = IasZone.ZoneStatus.Alarm_1
+        else:
+            migrated_state = IasZone.ZoneStatus(0)
+
+        self._channel.cluster.update_attribute(
+            IasZone.attributes_by_name[self.SENSOR_ATTR].id, migrated_state
+        )
+
 
 @MULTI_MATCH(
     channel_names="tuya_manufacturer",

tests/components/zha/test_binary_sensor.py

@@ -8,12 +8,15 @@ import zigpy.zcl.clusters.security as security
 
 from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE, Platform
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import restore_state
+from homeassistant.util import dt as dt_util
 
 from .common import (
     async_enable_traffic,
     async_test_rejoin,
     find_entity_id,
     send_attributes_report,
+    update_attribute_cache,
 )
 from .conftest import SIG_EP_INPUT, SIG_EP_OUTPUT, SIG_EP_PROFILE, SIG_EP_TYPE
 
@@ -120,3 +123,92 @@ async def test_binary_sensor(
     # test rejoin
     await async_test_rejoin(hass, zigpy_device, [cluster], reporting)
     assert hass.states.get(entity_id).state == STATE_OFF
+
+
+@pytest.fixture
+def core_rs(hass_storage):
+    """Core.restore_state fixture."""
+
+    def _storage(entity_id, attributes, state):
+        now = dt_util.utcnow().isoformat()
+
+        hass_storage[restore_state.STORAGE_KEY] = {
+            "version": restore_state.STORAGE_VERSION,
+            "key": restore_state.STORAGE_KEY,
+            "data": [
+                {
+                    "state": {
+                        "entity_id": entity_id,
+                        "state": str(state),
+                        "attributes": attributes,
+                        "last_changed": now,
+                        "last_updated": now,
+                        "context": {
+                            "id": "3c2243ff5f30447eb12e7348cfd5b8ff",
+                            "user_id": None,
+                        },
+                    },
+                    "last_seen": now,
+                }
+            ],
+        }
+        return
+
+    return _storage
+
+
+@pytest.mark.parametrize(
+    "restored_state",
+    [
+        STATE_ON,
+        STATE_OFF,
+    ],
+)
+async def test_binary_sensor_migration_not_migrated(
+    hass: HomeAssistant,
+    zigpy_device_mock,
+    core_rs,
+    zha_device_restored,
+    restored_state,
+) -> None:
+    """Test temporary ZHA IasZone binary_sensor migration to zigpy cache."""
+
+    entity_id = "binary_sensor.fakemanufacturer_fakemodel_iaszone"
+    core_rs(entity_id, state=restored_state, attributes={})  # migration sensor state
+
+    zigpy_device = zigpy_device_mock(DEVICE_IAS)
+    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
+
+    # confirm migration extra state attribute was set to True
+    assert hass.states.get(entity_id).attributes["migrated_to_cache"]
+
+
+async def test_binary_sensor_migration_already_migrated(
+    hass: HomeAssistant,
+    zigpy_device_mock,
+    core_rs,
+    zha_device_restored,
+) -> None:
+    """Test temporary ZHA IasZone binary_sensor migration doesn't migrate multiple times."""
+
+    entity_id = "binary_sensor.fakemanufacturer_fakemodel_iaszone"
+    core_rs(entity_id, state=STATE_OFF, attributes={"migrated_to_cache": True})
+
+    zigpy_device = zigpy_device_mock(DEVICE_IAS)
+
+    cluster = zigpy_device.endpoints.get(1).ias_zone
+    cluster.PLUGGED_ATTR_READS = {
+        "zone_status": security.IasZone.ZoneStatus.Alarm_1,
+    }
+    update_attribute_cache(cluster)
+
+    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 == STATE_ON  # matches attribute cache
+    assert hass.states.get(entity_id).attributes["migrated_to_cache"]