Set up single coordinator for all config entries in IronOS (#129108)

homeassistant/components/iron_os/__init__.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
 import logging
 from typing import TYPE_CHECKING
 
@@ -14,7 +13,10 @@ from homeassistant.config_entries import ConfigEntry
 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
+from homeassistant.util.hass_dict import HassKey
 
 from .const import DOMAIN
 from .coordinator import IronOSFirmwareUpdateCoordinator, IronOSLiveDataCoordinator
@@ -22,19 +24,25 @@ from .coordinator import IronOSFirmwareUpdateCoordinator, IronOSLiveDataCoordina
 PLATFORMS: list[Platform] = [Platform.NUMBER, Platform.SENSOR, Platform.UPDATE]
 
 
-@dataclass
+type IronOSConfigEntry = ConfigEntry[IronOSLiveDataCoordinator]
-class IronOSCoordinators:
+IRON_OS_KEY: HassKey[IronOSFirmwareUpdateCoordinator] = HassKey(DOMAIN)
-    """IronOS data class holding coordinators."""
 
-    live_data: IronOSLiveDataCoordinator
+CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
-    firmware: IronOSFirmwareUpdateCoordinator
-
-
-type IronOSConfigEntry = ConfigEntry[IronOSCoordinators]
 
 _LOGGER = logging.getLogger(__name__)
 
 
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up IronOS firmware update coordinator."""
+
+    session = async_get_clientsession(hass)
+    github = GitHubAPI(session=session)
+
+    hass.data[IRON_OS_KEY] = IronOSFirmwareUpdateCoordinator(hass, github)
+    await hass.data[IRON_OS_KEY].async_request_refresh()
+    return True
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: IronOSConfigEntry) -> bool:
     """Set up IronOS from a config entry."""
     if TYPE_CHECKING:
@@ -54,16 +62,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: IronOSConfigEntry) -> bo
     coordinator = IronOSLiveDataCoordinator(hass, device)
     await coordinator.async_config_entry_first_refresh()
 
-    session = async_get_clientsession(hass)
+    entry.runtime_data = coordinator
-    github = GitHubAPI(session=session)
-
-    firmware_update_coordinator = IronOSFirmwareUpdateCoordinator(hass, device, github)
-    await firmware_update_coordinator.async_config_entry_first_refresh()
-
-    entry.runtime_data = IronOSCoordinators(
-        live_data=coordinator,
-        firmware=firmware_update_coordinator,
-    )
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True

homeassistant/components/iron_os/coordinator.py

@@ -21,24 +21,19 @@ SCAN_INTERVAL = timedelta(seconds=5)
 SCAN_INTERVAL_GITHUB = timedelta(hours=3)
 
 
-class IronOSBaseCoordinator[_DataT](DataUpdateCoordinator[_DataT]):
+class IronOSLiveDataCoordinator(DataUpdateCoordinator[LiveDataResponse]):
-    """IronOS base coordinator."""
+    """IronOS live data coordinator."""
 
     device_info: DeviceInfoResponse
     config_entry: ConfigEntry
 
-    def __init__(
+    def __init__(self, hass: HomeAssistant, device: Pynecil) -> None:
-        self,
-        hass: HomeAssistant,
-        device: Pynecil,
-        update_interval: timedelta,
-    ) -> None:
         """Initialize IronOS coordinator."""
         super().__init__(
             hass,
             _LOGGER,
             name=DOMAIN,
-            update_interval=update_interval,
+            update_interval=SCAN_INTERVAL,
         )
         self.device = device
 
@@ -47,14 +42,6 @@ class IronOSBaseCoordinator[_DataT](DataUpdateCoordinator[_DataT]):
 
         self.device_info = await self.device.get_device_info()
 
-
-class IronOSLiveDataCoordinator(IronOSBaseCoordinator):
-    """IronOS live data coordinator."""
-
-    def __init__(self, hass: HomeAssistant, device: Pynecil) -> None:
-        """Initialize IronOS coordinator."""
-        super().__init__(hass, device=device, update_interval=SCAN_INTERVAL)
-
     async def _async_update_data(self) -> LiveDataResponse:
         """Fetch data from Device."""
 
@@ -65,12 +52,17 @@ class IronOSLiveDataCoordinator(IronOSBaseCoordinator):
             raise UpdateFailed("Cannot connect to device") from e
 
 
-class IronOSFirmwareUpdateCoordinator(IronOSBaseCoordinator):
+class IronOSFirmwareUpdateCoordinator(DataUpdateCoordinator[GitHubReleaseModel]):
     """IronOS coordinator for retrieving update information from github."""
 
-    def __init__(self, hass: HomeAssistant, device: Pynecil, github: GitHubAPI) -> None:
+    def __init__(self, hass: HomeAssistant, github: GitHubAPI) -> None:
         """Initialize IronOS coordinator."""
-        super().__init__(hass, device=device, update_interval=SCAN_INTERVAL_GITHUB)
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=SCAN_INTERVAL_GITHUB,
+        )
         self.github = github
 
     async def _async_update_data(self) -> GitHubReleaseModel:

homeassistant/components/iron_os/entity.py

@@ -9,17 +9,17 @@ from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import MANUFACTURER, MODEL
-from .coordinator import IronOSBaseCoordinator
+from .coordinator import IronOSLiveDataCoordinator
 
 
-class IronOSBaseEntity(CoordinatorEntity[IronOSBaseCoordinator]):
+class IronOSBaseEntity(CoordinatorEntity[IronOSLiveDataCoordinator]):
     """Base IronOS entity."""
 
     _attr_has_entity_name = True
 
     def __init__(
         self,
-        coordinator: IronOSBaseCoordinator,
+        coordinator: IronOSLiveDataCoordinator,
         entity_description: EntityDescription,
     ) -> None:
         """Initialize the sensor."""

homeassistant/components/iron_os/number.py

@@ -61,7 +61,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up number entities from a config entry."""
-    coordinator = entry.runtime_data.live_data
+    coordinator = entry.runtime_data
 
     async_add_entities(
         IronOSNumberEntity(coordinator, description)

homeassistant/components/iron_os/sensor.py

@@ -180,7 +180,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up sensors from a config entry."""
-    coordinator = entry.runtime_data.live_data
+    coordinator = entry.runtime_data
 
     async_add_entities(
         IronOSSensorEntity(coordinator, description)

homeassistant/components/iron_os/update.py

@@ -11,8 +11,8 @@ from homeassistant.components.update import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import IronOSConfigEntry
+from . import IRON_OS_KEY, IronOSConfigEntry, IronOSLiveDataCoordinator
-from .coordinator import IronOSBaseCoordinator
+from .coordinator import IronOSFirmwareUpdateCoordinator
 from .entity import IronOSBaseEntity
 
 UPDATE_DESCRIPTION = UpdateEntityDescription(
@@ -28,9 +28,11 @@ async def async_setup_entry(
 ) -> None:
     """Set up IronOS update platform."""
 
-    coordinator = entry.runtime_data.firmware
+    coordinator = entry.runtime_data
 
-    async_add_entities([IronOSUpdate(coordinator, UPDATE_DESCRIPTION)])
+    async_add_entities(
+        [IronOSUpdate(coordinator, hass.data[IRON_OS_KEY], UPDATE_DESCRIPTION)]
+    )
 
 
 class IronOSUpdate(IronOSBaseEntity, UpdateEntity):
@@ -40,10 +42,12 @@ class IronOSUpdate(IronOSBaseEntity, UpdateEntity):
 
     def __init__(
         self,
-        coordinator: IronOSBaseCoordinator,
+        coordinator: IronOSLiveDataCoordinator,
+        firmware_update: IronOSFirmwareUpdateCoordinator,
         entity_description: UpdateEntityDescription,
     ) -> None:
         """Initialize the sensor."""
+        self.firmware_update = firmware_update
         super().__init__(coordinator, entity_description)
 
     @property
@@ -56,21 +60,36 @@ class IronOSUpdate(IronOSBaseEntity, UpdateEntity):
     def title(self) -> str | None:
         """Title of the IronOS release."""
 
-        return f"IronOS {self.coordinator.data.name}"
+        return f"IronOS {self.firmware_update.data.name}"
 
     @property
     def release_url(self) -> str | None:
         """URL to the full release notes of the latest IronOS version available."""
 
-        return self.coordinator.data.html_url
+        return self.firmware_update.data.html_url
 
     @property
     def latest_version(self) -> str | None:
         """Latest IronOS version available for install."""
 
-        return self.coordinator.data.tag_name
+        return self.firmware_update.data.tag_name
 
     async def async_release_notes(self) -> str | None:
         """Return the release notes."""
 
-        return self.coordinator.data.body
+        return self.firmware_update.data.body
+
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass.
+
+        Register extra update listener for the firmware update coordinator.
+        """
+        await super().async_added_to_hass()
+        self.async_on_remove(
+            self.firmware_update.async_add_listener(self._handle_coordinator_update)
+        )
+
+    @property
+    def available(self) -> bool:
+        """Return if entity is available."""
+        return super().available and self.firmware_update.last_update_success

tests/components/iron_os/test_update.py

@@ -8,7 +8,7 @@ import pytest
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.config_entries import ConfigEntryState
-from homeassistant.const import Platform
+from homeassistant.const import STATE_UNAVAILABLE, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
@@ -57,12 +57,12 @@ async def test_update(
 
 
 @pytest.mark.usefixtures("ble_device", "mock_pynecil")
-async def test_config_entry_not_ready(
+async def test_update_unavailable(
     hass: HomeAssistant,
     config_entry: MockConfigEntry,
     mock_githubapi: AsyncMock,
 ) -> None:
-    """Test config entry not ready."""
+    """Test update entity unavailable on error."""
 
     mock_githubapi.repos.releases.latest.side_effect = GitHubException
 
@@ -70,4 +70,8 @@ async def test_config_entry_not_ready(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert config_entry.state is ConfigEntryState.SETUP_RETRY
+    assert config_entry.state is ConfigEntryState.LOADED
+
+    state = hass.states.get("update.pinecil_firmware")
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE