From 6c6edba4ecb258a0f741e1cc089e9078a0b1a55e Mon Sep 17 00:00:00 2001
From: Manu <4445816+tr4nt0r@users.noreply.github.com>
Date: Wed, 18 Dec 2024 23:21:37 +0000
Subject: [PATCH] use restore_cache helper

---
 tests/components/iron_os/test_update.py | 34 +++++++++++++------------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/tests/components/iron_os/test_update.py b/tests/components/iron_os/test_update.py
index 10417d5e81221..137d42a5d51c9 100644
--- a/tests/components/iron_os/test_update.py
+++ b/tests/components/iron_os/test_update.py
@@ -13,7 +13,7 @@
 from homeassistant.core import HomeAssistant, State
 from homeassistant.helpers import entity_registry as er
 
-from tests.common import MockConfigEntry, snapshot_platform
+from tests.common import MockConfigEntry, mock_restore_cache, snapshot_platform
 from tests.typing import WebSocketGenerator
 
 
@@ -87,21 +87,23 @@ async def test_update_restore_last_state(
     """Test update entity restore last state."""
 
     mock_pynecil.get_device_info.side_effect = CommunicationError
-    with patch(
-        "homeassistant.components.iron_os.update.IronOSUpdate.async_get_last_state",
-        return_value=State(
-            "update.pinecil_firmware",
-            STATE_ON,
-            attributes={ATTR_INSTALLED_VERSION: "v2.21"},
+    mock_restore_cache(
+        hass,
+        (
+            State(
+                "update.pinecil_firmware",
+                STATE_ON,
+                attributes={ATTR_INSTALLED_VERSION: "v2.21"},
+            ),
         ),
-    ):
-        config_entry.add_to_hass(hass)
-        await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
+    )
+    config_entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
 
-        assert config_entry.state is ConfigEntryState.LOADED
+    assert config_entry.state is ConfigEntryState.LOADED
 
-        state = hass.states.get("update.pinecil_firmware")
-        assert state is not None
-        assert state.state == STATE_ON
-        assert state.attributes[ATTR_INSTALLED_VERSION] == "v2.21"
+    state = hass.states.get("update.pinecil_firmware")
+    assert state is not None
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_INSTALLED_VERSION] == "v2.21"