Fix zwave_js.update entity restore logic (#94043)

homeassistant/components/zwave_js/update.py

@@ -42,6 +42,7 @@ PARALLEL_UPDATES = 1
 
 UPDATE_DELAY_STRING = "delay"
 UPDATE_DELAY_INTERVAL = 5  # In minutes
+ATTR_LATEST_VERSION_FIRMWARE = "latest_version_firmware"
 
 
 @dataclass
@@ -53,7 +54,7 @@ class ZWaveNodeFirmwareUpdateExtraStoredData(ExtraStoredData):
     def as_dict(self) -> dict[str, Any]:
         """Return a dict representation of the extra data."""
         return {
-            "latest_version_firmware": asdict(self.latest_version_firmware)
+            ATTR_LATEST_VERSION_FIRMWARE: asdict(self.latest_version_firmware)
             if self.latest_version_firmware
             else None
         }
@@ -61,7 +62,7 @@ class ZWaveNodeFirmwareUpdateExtraStoredData(ExtraStoredData):
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> ZWaveNodeFirmwareUpdateExtraStoredData:
         """Initialize the extra data from a dict."""
-        if not (firmware_dict := data["latest_version_firmware"]):
+        if not (firmware_dict := data[ATTR_LATEST_VERSION_FIRMWARE]):
             return cls(None)
 
         return cls(NodeFirmwareUpdateInfo.from_dict(firmware_dict))
@@ -326,20 +327,24 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
         )
 
         # If we have a complete previous state, use that to set the latest version
-        if (state := await self.async_get_last_state()) and (
-            extra_data := await self.async_get_last_extra_data()
+        if (
+            (state := await self.async_get_last_state())
+            and (latest_version := state.attributes.get(ATTR_LATEST_VERSION))
+            is not None
+            and (extra_data := await self.async_get_last_extra_data())
         ):
-            self._attr_latest_version = state.attributes[ATTR_LATEST_VERSION]
+            self._attr_latest_version = latest_version
             self._latest_version_firmware = (
                 ZWaveNodeFirmwareUpdateExtraStoredData.from_dict(
                     extra_data.as_dict()
                 ).latest_version_firmware
             )
-        # If we have no state to restore, we can set the latest version to installed
-        # so that the entity starts as off. If we have partial restore data due to an
-        # upgrade to an HA version where this feature is released from one that is not
-        # the entity will start in an unknown state until we can correct on next update
-        elif not state:
+        # If we have no state or latest version to restore, we can set the latest
+        # version to installed so that the entity starts as off. If we have partial
+        # restore data due to an upgrade to an HA version where this feature is released
+        # from one that is not the entity will start in an unknown state until we can
+        # correct on next update
+        elif not state or not latest_version:
             self._attr_latest_version = self._attr_installed_version
 
         # Spread updates out in 5 minute increments to avoid flooding the network

tests/components/zwave_js/test_update.py

@@ -778,6 +778,42 @@ async def test_update_entity_full_restore_data_no_update_available(
     assert state.attributes[ATTR_LATEST_VERSION] == "10.7"
 
 
+async def test_update_entity_no_latest_version(
+    hass: HomeAssistant,
+    client,
+    climate_radio_thermostat_ct100_plus_different_endpoints,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test entity with no `latest_version` attr restores state."""
+    mock_restore_cache_with_extra_data(
+        hass,
+        [
+            (
+                State(
+                    UPDATE_ENTITY,
+                    STATE_OFF,
+                    {
+                        ATTR_INSTALLED_VERSION: "10.7",
+                        ATTR_LATEST_VERSION: None,
+                        ATTR_SKIPPED_VERSION: None,
+                    },
+                ),
+                {"latest_version_firmware": None},
+            )
+        ],
+    )
+    entry = MockConfigEntry(domain="zwave_js", data={"url": "ws://test.org"})
+    entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(UPDATE_ENTITY)
+    assert state
+    assert state.state == STATE_OFF
+    assert state.attributes[ATTR_SKIPPED_VERSION] is None
+    assert state.attributes[ATTR_LATEST_VERSION] == "10.7"
+
+
 async def test_update_entity_unload_asleep_node(
     hass: HomeAssistant, client, wallmote_central_scene, integration
 ) -> None: