Improve device handling for disconnected IronOS devices (#143446)

* Improve device handling for disconnected IronOS devices

* requested changes

* ble_device

homeassistant/components/iron_os/__init__.py

@@ -7,10 +7,8 @@ from typing import TYPE_CHECKING
 
 from pynecil import IronOSUpdate, Pynecil
 
-from homeassistant.components import bluetooth
+from homeassistant.const import Platform
-from homeassistant.const import CONF_NAME, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.typing import ConfigType
@@ -35,7 +33,6 @@ PLATFORMS: list[Platform] = [
     Platform.UPDATE,
 ]
 
-
 CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
 
 
@@ -60,17 +57,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: IronOSConfigEntry) -> bo
     """Set up IronOS from a config entry."""
     if TYPE_CHECKING:
         assert entry.unique_id
-    ble_device = bluetooth.async_ble_device_from_address(
-        hass, entry.unique_id, connectable=True
-    )
-    if not ble_device:
-        raise ConfigEntryNotReady(
-            translation_domain=DOMAIN,
-            translation_key="setup_device_unavailable_exception",
-            translation_placeholders={CONF_NAME: entry.title},
-        )
 
-    device = Pynecil(ble_device)
+    device = Pynecil(entry.unique_id)
 
     live_data = IronOSLiveDataCoordinator(hass, entry, device)
     await live_data.async_config_entry_first_refresh()

homeassistant/components/iron_os/config_flow.py

@@ -2,9 +2,12 @@
 
 from __future__ import annotations
 
+import logging
 from typing import Any
 
+from bleak.exc import BleakError
 from habluetooth import BluetoothServiceInfoBleak
+from pynecil import CommunicationError, Pynecil
 import voluptuous as vol
 
 from homeassistant.components.bluetooth.api import async_discovered_service_info
@@ -13,6 +16,8 @@ from homeassistant.const import CONF_ADDRESS
 
 from .const import DISCOVERY_SVC_UUID, DOMAIN
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class IronOSConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for IronOS."""
@@ -36,30 +41,62 @@ class IronOSConfigFlow(ConfigFlow, domain=DOMAIN):
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Confirm discovery."""
+
+        errors: dict[str, str] = {}
+
         assert self._discovery_info is not None
         discovery_info = self._discovery_info
         title = discovery_info.name
 
         if user_input is not None:
-            return self.async_create_entry(title=title, data={})
+            device = Pynecil(discovery_info.address)
+            try:
+                await device.connect()
+            except (CommunicationError, BleakError, TimeoutError):
+                _LOGGER.debug("Cannot connect:", exc_info=True)
+                errors["base"] = "cannot_connect"
+            except Exception:
+                _LOGGER.exception("Unexpected exception:")
+                errors["base"] = "unknown"
+            else:
+                return self.async_create_entry(title=title, data={})
+            finally:
+                await device.disconnect()
 
         self._set_confirm_only()
         placeholders = {"name": title}
         self.context["title_placeholders"] = placeholders
         return self.async_show_form(
-            step_id="bluetooth_confirm", description_placeholders=placeholders
+            step_id="bluetooth_confirm",
+            description_placeholders=placeholders,
+            errors=errors,
         )
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Handle the user step to pick discovered device."""
+
+        errors: dict[str, str] = {}
+
         if user_input is not None:
             address = user_input[CONF_ADDRESS]
             title = self._discovered_devices[address]
             await self.async_set_unique_id(address, raise_on_progress=False)
             self._abort_if_unique_id_configured()
-            return self.async_create_entry(title=title, data={})
+            device = Pynecil(address)
+            try:
+                await device.connect()
+            except (CommunicationError, BleakError, TimeoutError):
+                _LOGGER.debug("Cannot connect:", exc_info=True)
+                errors["base"] = "cannot_connect"
+            except Exception:
+                _LOGGER.exception("Unexpected exception")
+                errors["base"] = "unknown"
+            else:
+                return self.async_create_entry(title=title, data={})
+            finally:
+                await device.disconnect()
 
         current_addresses = self._async_current_ids(include_ignore=False)
         for discovery_info in async_discovered_service_info(self.hass, True):
@@ -80,4 +117,5 @@ class IronOSConfigFlow(ConfigFlow, domain=DOMAIN):
             data_schema=vol.Schema(
                 {vol.Required(CONF_ADDRESS): vol.In(self._discovered_devices)}
             ),
+            errors=errors,
         )

homeassistant/components/iron_os/coordinator.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass
 from datetime import timedelta
 from enum import Enum
 import logging
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 from awesomeversion import AwesomeVersion
 from pynecil import (
@@ -22,10 +22,11 @@ from pynecil import (
 )
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers.debounce import Debouncer
+import homeassistant.helpers.device_registry as dr
+from homeassistant.helpers.device_registry import CONNECTION_BLUETOOTH
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN
@@ -83,14 +84,13 @@ class IronOSBaseCoordinator[_DataT](DataUpdateCoordinator[_DataT]):
         try:
             self.device_info = await self.device.get_device_info()
 
-        except CommunicationError as e:
+        except (CommunicationError, TimeoutError):
-            raise UpdateFailed(
+            self.device_info = DeviceInfoResponse()
-                translation_domain=DOMAIN,
-                translation_key="cannot_connect",
-                translation_placeholders={CONF_NAME: self.config_entry.title},
-            ) from e
 
-        self.v223_features = AwesomeVersion(self.device_info.build) >= V223
+        self.v223_features = (
+            self.device_info.build is not None
+            and AwesomeVersion(self.device_info.build) >= V223
+        )
 
 
 class IronOSLiveDataCoordinator(IronOSBaseCoordinator[LiveDataResponse]):
@@ -101,23 +101,18 @@ class IronOSLiveDataCoordinator(IronOSBaseCoordinator[LiveDataResponse]):
     ) -> None:
         """Initialize IronOS coordinator."""
         super().__init__(hass, config_entry, device, SCAN_INTERVAL)
+        self.device_info = DeviceInfoResponse()
 
     async def _async_update_data(self) -> LiveDataResponse:
         """Fetch data from Device."""
 
         try:
-            # device info is cached and won't be refetched on every
+            await self._update_device_info()
-            # coordinator refresh, only after the device has disconnected
-            # the device info is refetched
-            self.device_info = await self.device.get_device_info()
             return await self.device.get_live_data()
 
-        except CommunicationError as e:
+        except CommunicationError:
-            raise UpdateFailed(
+            _LOGGER.debug("Cannot connect to device", exc_info=True)
-                translation_domain=DOMAIN,
+            return self.data or LiveDataResponse()
-                translation_key="cannot_connect",
-                translation_placeholders={CONF_NAME: self.config_entry.title},
-            ) from e
 
     @property
     def has_tip(self) -> bool:
@@ -130,6 +125,32 @@ class IronOSLiveDataCoordinator(IronOSBaseCoordinator[LiveDataResponse]):
             return self.data.live_temp <= threshold
         return False
 
+    async def _update_device_info(self) -> None:
+        """Update device info.
+
+        device info is cached and won't be refetched on every
+        coordinator refresh, only after the device has disconnected
+        the device info is refetched.
+        """
+        build = self.device_info.build
+        self.device_info = await self.device.get_device_info()
+
+        if build == self.device_info.build:
+            return
+        device_registry = dr.async_get(self.hass)
+        if TYPE_CHECKING:
+            assert self.config_entry.unique_id
+        device = device_registry.async_get_device(
+            connections={(CONNECTION_BLUETOOTH, self.config_entry.unique_id)}
+        )
+        if device is None:
+            return
+        device_registry.async_update_device(
+            device_id=device.id,
+            sw_version=self.device_info.build,
+            serial_number=f"{self.device_info.device_sn} (ID:{self.device_info.device_id})",
+        )
+
 
 class IronOSSettingsCoordinator(IronOSBaseCoordinator[SettingsDataResponse]):
     """IronOS coordinator."""

homeassistant/components/iron_os/entity.py

@@ -37,6 +37,16 @@ class IronOSBaseEntity(CoordinatorEntity[IronOSLiveDataCoordinator]):
             manufacturer=MANUFACTURER,
             model=MODEL,
             name="Pinecil",
-            sw_version=coordinator.device_info.build,
-            serial_number=f"{coordinator.device_info.device_sn} (ID:{coordinator.device_info.device_id})",
         )
+        if coordinator.device_info.is_synced:
+            self._attr_device_info.update(
+                DeviceInfo(
+                    sw_version=coordinator.device_info.build,
+                    serial_number=f"{coordinator.device_info.device_sn} (ID:{coordinator.device_info.device_id})",
+                )
+            )
+
+    @property
+    def available(self) -> bool:
+        """Return if entity is available."""
+        return super().available and self.coordinator.device.is_connected

homeassistant/components/iron_os/quality_scale.yaml

@@ -21,10 +21,10 @@ rules:
   entity-unique-id: done
   has-entity-name: done
   runtime-data: done
-  test-before-configure:
+  test-before-configure: done
+  test-before-setup:
     status: exempt
-    comment: Device is set up from a Bluetooth discovery
+    comment: Device is expected to be disconnected most of the time but will connect quickly when reachable
-  test-before-setup: done
   unique-config-entry: done
 
   # Silver
@@ -47,8 +47,8 @@ rules:
   devices: done
   diagnostics: done
   discovery-update-info:
-    status: exempt
+    status: done
-    comment: Device is not connected to an ip network. Other information from discovery is immutable and does not require updating.
+    comment: Device is not connected to an ip network. FW version in device info is updated.
   discovery: done
   docs-data-update: done
   docs-examples: done

homeassistant/components/iron_os/strings.json

@@ -20,7 +20,13 @@
     },
     "abort": {
       "no_devices_found": "[%key:common::config_flow::abort::no_devices_found%]",
-      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]"
+      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
+      "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
+      "unknown": "[%key:common::config_flow::error::unknown%]"
+    },
+    "error": {
+      "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
+      "unknown": "[%key:common::config_flow::error::unknown%]"
     }
   },
   "entity": {
@@ -276,12 +282,6 @@
     }
   },
   "exceptions": {
-    "setup_device_unavailable_exception": {
-      "message": "Device {name} is not reachable"
-    },
-    "setup_device_connection_error_exception": {
-      "message": "Connection to device {name} failed, try again later"
-    },
     "submit_setting_failed": {
       "message": "Failed to submit setting to device, try again later"
     },

homeassistant/components/iron_os/update.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from homeassistant.components.update import (
+    ATTR_INSTALLED_VERSION,
     UpdateDeviceClass,
     UpdateEntity,
     UpdateEntityDescription,
@@ -10,6 +11,7 @@ from homeassistant.components.update import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
 
 from . import IRON_OS_KEY, IronOSConfigEntry, IronOSLiveDataCoordinator
 from .coordinator import IronOSFirmwareUpdateCoordinator
@@ -37,7 +39,7 @@ async def async_setup_entry(
     )
 
 
-class IronOSUpdate(IronOSBaseEntity, UpdateEntity):
+class IronOSUpdate(IronOSBaseEntity, UpdateEntity, RestoreEntity):
     """Representation of an IronOS update entity."""
 
     _attr_supported_features = UpdateEntityFeature.RELEASE_NOTES
@@ -56,7 +58,7 @@ class IronOSUpdate(IronOSBaseEntity, UpdateEntity):
     def installed_version(self) -> str | None:
         """IronOS version on the device."""
 
-        return self.coordinator.device_info.build
+        return self.coordinator.device_info.build or self._attr_installed_version
 
     @property
     def title(self) -> str | None:
@@ -86,6 +88,9 @@ class IronOSUpdate(IronOSBaseEntity, UpdateEntity):
 
         Register extra update listener for the firmware update coordinator.
         """
+        if state := await self.async_get_last_state():
+            self._attr_installed_version = state.attributes.get(ATTR_INSTALLED_VERSION)
+
         await super().async_added_to_hass()
         self.async_on_remove(
             self.firmware_update.async_add_listener(self._handle_coordinator_update)

tests/components/iron_os/conftest.py

@@ -159,9 +159,10 @@ def mock_ironosupdate() -> Generator[AsyncMock]:
 @pytest.fixture
 def mock_pynecil() -> Generator[AsyncMock]:
     """Mock Pynecil library."""
-    with patch(
+    with (
-        "homeassistant.components.iron_os.Pynecil", autospec=True
+        patch("homeassistant.components.iron_os.Pynecil", autospec=True) as mock_client,
-    ) as mock_client:
+        patch("homeassistant.components.iron_os.config_flow.Pynecil", new=mock_client),
+    ):
         client = mock_client.return_value
 
         client.get_device_info.return_value = DeviceInfoResponse(
@@ -170,6 +171,7 @@ def mock_pynecil() -> Generator[AsyncMock]:
             address="c0:ff:ee:c0:ff:ee",
             device_sn="0000c0ffeec0ffee",
             name=DEFAULT_NAME,
+            is_synced=True,
         )
         client.get_settings.return_value = SettingsDataResponse(
             sleep_temp=150,
@@ -225,4 +227,6 @@ def mock_pynecil() -> Generator[AsyncMock]:
             operating_mode=OperatingMode.SOLDERING,
             estimated_power=24.8,
         )
+        client._client = AsyncMock()
+        client._client.return_value.is_connected = True
         yield client

tests/components/iron_os/snapshots/test_diagnostics.ambr

@@ -6,7 +6,7 @@
     }),
     'device_info': dict({
       '__type': "<class 'pynecil.types.DeviceInfoResponse'>",
-      'repr': "DeviceInfoResponse(build='v2.23', device_id='c0ffeeC0', address='c0:ff:ee:c0:ff:ee', device_sn='0000c0ffeec0ffee', name='Pinecil-C0FFEEE', is_synced=False)",
+      'repr': "DeviceInfoResponse(build='v2.23', device_id='c0ffeeC0', address='c0:ff:ee:c0:ff:ee', device_sn='0000c0ffeec0ffee', name='Pinecil-C0FFEEE', is_synced=True)",
     }),
     'live_data': dict({
       '__type': "<class 'pynecil.types.LiveDataResponse'>",

tests/components/iron_os/test_config_flow.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 from unittest.mock import AsyncMock, MagicMock
 
+from pynecil import CommunicationError
 import pytest
 
 from homeassistant.components.iron_os import DOMAIN
@@ -16,7 +17,7 @@ from .conftest import DEFAULT_NAME, PINECIL_SERVICE_INFO, USER_INPUT
 from tests.common import MockConfigEntry
 
 
-@pytest.mark.usefixtures("discovery")
+@pytest.mark.usefixtures("discovery", "mock_pynecil")
 async def test_async_step_user(
     hass: HomeAssistant, mock_setup_entry: AsyncMock
 ) -> None:
@@ -34,10 +35,52 @@ async def test_async_step_user(
     assert result["type"] is FlowResultType.CREATE_ENTRY
     assert result["title"] == DEFAULT_NAME
     assert result["data"] == {}
+    assert result["result"].unique_id == "c0:ff:ee:c0:ff:ee"
     assert len(mock_setup_entry.mock_calls) == 1
 
 
+@pytest.mark.parametrize(
+    ("raise_error", "text_error"),
+    [
+        (CommunicationError, "cannot_connect"),
+        (Exception, "unknown"),
+    ],
+)
 @pytest.mark.usefixtures("discovery")
+async def test_async_step_user_errors(
+    hass: HomeAssistant,
+    mock_pynecil: AsyncMock,
+    raise_error: Exception,
+    text_error: str,
+) -> None:
+    """Test the user config flow errors."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": SOURCE_USER}
+    )
+    assert result["type"] is FlowResultType.FORM
+
+    mock_pynecil.connect.side_effect = raise_error
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        USER_INPUT,
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["errors"] == {"base": text_error}
+
+    mock_pynecil.connect.side_effect = None
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        USER_INPUT,
+    )
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["title"] == DEFAULT_NAME
+    assert result["data"] == {}
+    assert result["result"].unique_id == "c0:ff:ee:c0:ff:ee"
+
+
+@pytest.mark.usefixtures("discovery", "mock_pynecil")
 async def test_async_step_user_device_added_between_steps(
     hass: HomeAssistant, config_entry: MockConfigEntry
 ) -> None:
@@ -73,6 +116,7 @@ async def test_form_no_device_discovered(
     assert result["reason"] == "no_devices_found"
 
 
+@pytest.mark.usefixtures("mock_pynecil")
 async def test_async_step_bluetooth(hass: HomeAssistant) -> None:
     """Test discovery via bluetooth."""
     result = await hass.config_entries.flow.async_init(
@@ -92,6 +136,49 @@ async def test_async_step_bluetooth(hass: HomeAssistant) -> None:
     assert result["result"].unique_id == "c0:ff:ee:c0:ff:ee"
 
 
+@pytest.mark.parametrize(
+    ("raise_error", "text_error"),
+    [
+        (CommunicationError, "cannot_connect"),
+        (Exception, "unknown"),
+    ],
+)
+async def test_async_step_bluetooth_errors(
+    hass: HomeAssistant,
+    mock_pynecil: AsyncMock,
+    raise_error: Exception,
+    text_error: str,
+) -> None:
+    """Test discovery via bluetooth errors."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": SOURCE_BLUETOOTH},
+        data=PINECIL_SERVICE_INFO,
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "bluetooth_confirm"
+
+    mock_pynecil.connect.side_effect = raise_error
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], user_input={}
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["errors"] == {"base": text_error}
+
+    mock_pynecil.connect.side_effect = None
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        USER_INPUT,
+    )
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["title"] == DEFAULT_NAME
+    assert result["data"] == {}
+    assert result["result"].unique_id == "c0:ff:ee:c0:ff:ee"
+
+
+@pytest.mark.usefixtures("mock_pynecil")
 async def test_async_step_bluetooth_devices_already_setup(
     hass: HomeAssistant, config_entry: AsyncMock
 ) -> None:
@@ -108,7 +195,7 @@ async def test_async_step_bluetooth_devices_already_setup(
     assert result["reason"] == "already_configured"
 
 
-@pytest.mark.usefixtures("discovery")
+@pytest.mark.usefixtures("discovery", "mock_pynecil")
 async def test_async_step_user_setup_replaces_igonored_device(
     hass: HomeAssistant, config_entry_ignored: AsyncMock
 ) -> None:

tests/components/iron_os/test_init.py

@@ -10,6 +10,8 @@ import pytest
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import STATE_UNKNOWN
 from homeassistant.core import HomeAssistant
+import homeassistant.helpers.device_registry as dr
+from homeassistant.helpers.device_registry import CONNECTION_BLUETOOTH
 
 from .conftest import DEFAULT_NAME
 
@@ -35,41 +37,6 @@ async def test_setup_and_unload(
     assert config_entry.state is ConfigEntryState.NOT_LOADED
 
 
-@pytest.mark.usefixtures("ble_device")
-async def test_update_data_config_entry_not_ready(
-    hass: HomeAssistant,
-    config_entry: MockConfigEntry,
-    mock_pynecil: AsyncMock,
-) -> None:
-    """Test config entry not ready."""
-    mock_pynecil.get_live_data.side_effect = CommunicationError
-    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.SETUP_RETRY
-
-
-@pytest.mark.usefixtures("entity_registry_enabled_by_default", "ble_device")
-async def test_setup_config_entry_not_ready(
-    hass: HomeAssistant,
-    config_entry: MockConfigEntry,
-    mock_pynecil: AsyncMock,
-    freezer: FrozenDateTimeFactory,
-) -> None:
-    """Test config entry not ready."""
-    mock_pynecil.get_settings.side_effect = CommunicationError
-    mock_pynecil.get_device_info.side_effect = CommunicationError
-    config_entry.add_to_hass(hass)
-    await hass.config_entries.async_setup(config_entry.entry_id)
-    await hass.async_block_till_done()
-    freezer.tick(timedelta(seconds=3))
-    async_fire_time_changed(hass)
-    await hass.async_block_till_done()
-
-    assert config_entry.state is ConfigEntryState.SETUP_RETRY
-
-
 @pytest.mark.usefixtures("entity_registry_enabled_by_default", "ble_device")
 async def test_settings_exception(
     hass: HomeAssistant,
@@ -123,3 +90,47 @@ async def test_v223_entities_not_loaded(
     ) is not None
 
     assert len(state.attributes["options"]) == 2
+
+
+@pytest.mark.usefixtures("ble_device")
+async def test_device_info_update(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    mock_pynecil: AsyncMock,
+    device_registry: dr.DeviceRegistry,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test device info gets updated."""
+
+    mock_pynecil.get_device_info.return_value = DeviceInfoResponse()
+    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
+
+    device = device_registry.async_get_device(
+        connections={(CONNECTION_BLUETOOTH, config_entry.unique_id)}
+    )
+    assert device
+    assert device.sw_version is None
+    assert device.serial_number is None
+
+    mock_pynecil.get_device_info.return_value = DeviceInfoResponse(
+        build="v2.22",
+        device_id="c0ffeeC0",
+        address="c0:ff:ee:c0:ff:ee",
+        device_sn="0000c0ffeec0ffee",
+        name=DEFAULT_NAME,
+    )
+
+    freezer.tick(timedelta(seconds=60))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    device = device_registry.async_get_device(
+        connections={(CONNECTION_BLUETOOTH, config_entry.unique_id)}
+    )
+    assert device
+    assert device.sw_version == "v2.22"
+    assert device.serial_number == "0000c0ffeec0ffee (ID:c0ffeeC0)"

tests/components/iron_os/test_sensor.py

@@ -4,7 +4,7 @@ from collections.abc import AsyncGenerator
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from freezegun.api import FrozenDateTimeFactory
-from pynecil import CommunicationError, LiveDataResponse
+from pynecil import LiveDataResponse
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
@@ -62,7 +62,7 @@ async def test_sensors_unavailable(
 
     assert config_entry.state is ConfigEntryState.LOADED
 
-    mock_pynecil.get_live_data.side_effect = CommunicationError
+    mock_pynecil.is_connected = False
     freezer.tick(SCAN_INTERVAL)
     async_fire_time_changed(hass)
 

tests/components/iron_os/test_update.py

@@ -3,16 +3,17 @@
 from collections.abc import AsyncGenerator
 from unittest.mock import AsyncMock, patch
 
-from pynecil import UpdateException
+from pynecil import CommunicationError, UpdateException
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
+from homeassistant.components.update import ATTR_INSTALLED_VERSION
 from homeassistant.config_entries import ConfigEntryState
-from homeassistant.const import STATE_UNAVAILABLE, Platform
+from homeassistant.const import STATE_ON, STATE_UNAVAILABLE, Platform
-from homeassistant.core import HomeAssistant
+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
 
 
@@ -75,3 +76,34 @@ async def test_update_unavailable(
     state = hass.states.get("update.pinecil_firmware")
     assert state is not None
     assert state.state == STATE_UNAVAILABLE
+
+
+@pytest.mark.usefixtures("ble_device")
+async def test_update_restore_last_state(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    mock_pynecil: AsyncMock,
+) -> None:
+    """Test update entity restore last state."""
+
+    mock_pynecil.get_device_info.side_effect = CommunicationError
+    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()
+
+    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"