From 2b3d57859ee994e84f9b1474782662d78d457879 Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Sat, 4 Nov 2023 08:12:06 -0400
Subject: [PATCH] Add test for firmware update scenario (#103314)

---
 homeassistant/components/zwave_js/update.py |  4 +++
 tests/components/zwave_js/test_update.py    | 36 +++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/homeassistant/components/zwave_js/update.py b/homeassistant/components/zwave_js/update.py
index e49eb8a2017f74..37cfdc68569128 100644
--- a/homeassistant/components/zwave_js/update.py
+++ b/homeassistant/components/zwave_js/update.py
@@ -333,6 +333,10 @@ async def async_added_to_hass(self) -> None:
             )
         )
 
+        # Make sure these variables are set for the elif evaluation
+        state = None
+        latest_version = None
+
         # If we have a complete previous state, use that to set the latest version
         if (
             (state := await self.async_get_last_state())
diff --git a/tests/components/zwave_js/test_update.py b/tests/components/zwave_js/test_update.py
index 46dca7a35eca20..9e17f25c708720 100644
--- a/tests/components/zwave_js/test_update.py
+++ b/tests/components/zwave_js/test_update.py
@@ -694,6 +694,42 @@ async def test_update_entity_partial_restore_data(
     assert state.state == STATE_UNKNOWN
 
 
+async def test_update_entity_partial_restore_data_2(
+    hass: HomeAssistant,
+    client,
+    climate_radio_thermostat_ct100_plus_different_endpoints,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test second scenario where update entity has partial restore data."""
+    mock_restore_cache_with_extra_data(
+        hass,
+        [
+            (
+                State(
+                    UPDATE_ENTITY,
+                    STATE_ON,
+                    {
+                        ATTR_INSTALLED_VERSION: "10.7",
+                        ATTR_LATEST_VERSION: "10.8",
+                        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_UNKNOWN
+    assert state.attributes[ATTR_SKIPPED_VERSION] is None
+    assert state.attributes[ATTR_LATEST_VERSION] is None
+
+
 async def test_update_entity_full_restore_data_skipped_version(
     hass: HomeAssistant,
     client,