From 78116f15960345ebe545048a3bd739a920affafb Mon Sep 17 00:00:00 2001 From: Manu <4445816+tr4nt0r@users.noreply.github.com> Date: Fri, 25 Oct 2024 10:51:23 +0200 Subject: [PATCH] Set up single coordinator for all config entries in IronOS (#129108) --- homeassistant/components/iron_os/__init__.py | 37 +++++++++---------- .../components/iron_os/coordinator.py | 32 ++++++---------- homeassistant/components/iron_os/entity.py | 6 +-- homeassistant/components/iron_os/number.py | 2 +- homeassistant/components/iron_os/sensor.py | 2 +- homeassistant/components/iron_os/update.py | 37 ++++++++++++++----- tests/components/iron_os/test_update.py | 12 ++++-- 7 files changed, 71 insertions(+), 57 deletions(-) diff --git a/homeassistant/components/iron_os/__init__.py b/homeassistant/components/iron_os/__init__.py index 43691c8594a..56a83117e68 100644 --- a/homeassistant/components/iron_os/__init__.py +++ b/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 -class IronOSCoordinators: - """IronOS data class holding coordinators.""" +type IronOSConfigEntry = ConfigEntry[IronOSLiveDataCoordinator] +IRON_OS_KEY: HassKey[IronOSFirmwareUpdateCoordinator] = HassKey(DOMAIN) - live_data: IronOSLiveDataCoordinator - firmware: IronOSFirmwareUpdateCoordinator - - -type IronOSConfigEntry = ConfigEntry[IronOSCoordinators] +CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) _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) - 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, - ) + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True diff --git a/homeassistant/components/iron_os/coordinator.py b/homeassistant/components/iron_os/coordinator.py index 175de484870..da82b76f92e 100644 --- a/homeassistant/components/iron_os/coordinator.py +++ b/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]): - """IronOS base coordinator.""" +class IronOSLiveDataCoordinator(DataUpdateCoordinator[LiveDataResponse]): + """IronOS live data coordinator.""" device_info: DeviceInfoResponse config_entry: ConfigEntry - def __init__( - self, - hass: HomeAssistant, - device: Pynecil, - update_interval: timedelta, - ) -> None: + def __init__(self, hass: HomeAssistant, device: Pynecil) -> 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: diff --git a/homeassistant/components/iron_os/entity.py b/homeassistant/components/iron_os/entity.py index d1c9a9aa0ee..77bebda9390 100644 --- a/homeassistant/components/iron_os/entity.py +++ b/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.""" diff --git a/homeassistant/components/iron_os/number.py b/homeassistant/components/iron_os/number.py index bc8da968187..9230faec1f1 100644 --- a/homeassistant/components/iron_os/number.py +++ b/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) diff --git a/homeassistant/components/iron_os/sensor.py b/homeassistant/components/iron_os/sensor.py index a44e61c4de3..095ffd254df 100644 --- a/homeassistant/components/iron_os/sensor.py +++ b/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) diff --git a/homeassistant/components/iron_os/update.py b/homeassistant/components/iron_os/update.py index 9086dc0b7b5..bae9ccd4c6c 100644 --- a/homeassistant/components/iron_os/update.py +++ b/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 .coordinator import IronOSBaseCoordinator +from . import IRON_OS_KEY, IronOSConfigEntry, IronOSLiveDataCoordinator +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 diff --git a/tests/components/iron_os/test_update.py b/tests/components/iron_os/test_update.py index 70336e69620..7a2650ba7a3 100644 --- a/tests/components/iron_os/test_update.py +++ b/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