From 5b91ea45502d7dfc36e71dd5c1dacb007a95de6b Mon Sep 17 00:00:00 2001 From: Robert Hillis Date: Wed, 12 Jun 2024 10:52:18 -0400 Subject: [PATCH] Store runtime data inside the config entry in Goalzero (#119440) --- homeassistant/components/goalzero/__init__.py | 17 ++++++----------- .../components/goalzero/binary_sensor.py | 12 +++++------- .../components/goalzero/coordinator.py | 4 +++- homeassistant/components/goalzero/sensor.py | 13 +++++-------- homeassistant/components/goalzero/switch.py | 13 +++++-------- tests/components/goalzero/__init__.py | 3 +-- 6 files changed, 25 insertions(+), 37 deletions(-) diff --git a/homeassistant/components/goalzero/__init__.py b/homeassistant/components/goalzero/__init__.py index 60b0338c258..6698d1efc99 100644 --- a/homeassistant/components/goalzero/__init__.py +++ b/homeassistant/components/goalzero/__init__.py @@ -6,20 +6,18 @@ from typing import TYPE_CHECKING from goalzero import Yeti, exceptions -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_HOST, Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.device_registry import format_mac -from .const import DOMAIN -from .coordinator import GoalZeroDataUpdateCoordinator +from .coordinator import GoalZeroConfigEntry, GoalZeroDataUpdateCoordinator PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SWITCH] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: GoalZeroConfigEntry) -> bool: """Set up Goal Zero Yeti from a config entry.""" mac = entry.unique_id @@ -38,16 +36,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: except exceptions.ConnectError as ex: raise ConfigEntryNotReady(f"Failed to connect to device: {ex}") from ex - coordinator = GoalZeroDataUpdateCoordinator(hass, api) - await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = GoalZeroDataUpdateCoordinator(hass, api) + await entry.runtime_data.async_config_entry_first_refresh() await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: GoalZeroConfigEntry) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - hass.data[DOMAIN].pop(entry.entry_id) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/goalzero/binary_sensor.py b/homeassistant/components/goalzero/binary_sensor.py index eec8773db30..6bd061879eb 100644 --- a/homeassistant/components/goalzero/binary_sensor.py +++ b/homeassistant/components/goalzero/binary_sensor.py @@ -9,12 +9,11 @@ from homeassistant.components.binary_sensor import ( BinarySensorEntity, BinarySensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .coordinator import GoalZeroConfigEntry from .entity import GoalZeroEntity PARALLEL_UPDATES = 0 @@ -43,14 +42,13 @@ BINARY_SENSOR_TYPES: tuple[BinarySensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: GoalZeroConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Goal Zero Yeti sensor.""" async_add_entities( - GoalZeroBinarySensor( - hass.data[DOMAIN][entry.entry_id], - description, - ) + GoalZeroBinarySensor(entry.runtime_data, description) for description in BINARY_SENSOR_TYPES ) diff --git a/homeassistant/components/goalzero/coordinator.py b/homeassistant/components/goalzero/coordinator.py index 61c3a8dba29..3c7cd967482 100644 --- a/homeassistant/components/goalzero/coordinator.py +++ b/homeassistant/components/goalzero/coordinator.py @@ -10,11 +10,13 @@ from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, Upda from .const import DOMAIN, LOGGER +type GoalZeroConfigEntry = ConfigEntry[GoalZeroDataUpdateCoordinator] + class GoalZeroDataUpdateCoordinator(DataUpdateCoordinator[None]): """Data update coordinator for the Goal zero integration.""" - config_entry: ConfigEntry + config_entry: GoalZeroConfigEntry def __init__(self, hass: HomeAssistant, api: Yeti) -> None: """Initialize the coordinator.""" diff --git a/homeassistant/components/goalzero/sensor.py b/homeassistant/components/goalzero/sensor.py index 86f8bc9455b..f565c216745 100644 --- a/homeassistant/components/goalzero/sensor.py +++ b/homeassistant/components/goalzero/sensor.py @@ -10,7 +10,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( PERCENTAGE, SIGNAL_STRENGTH_DECIBELS, @@ -26,7 +25,7 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType -from .const import DOMAIN +from .coordinator import GoalZeroConfigEntry from .entity import GoalZeroEntity SENSOR_TYPES: tuple[SensorEntityDescription, ...] = ( @@ -130,15 +129,13 @@ SENSOR_TYPES: tuple[SensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: GoalZeroConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Goal Zero Yeti sensor.""" async_add_entities( - GoalZeroSensor( - hass.data[DOMAIN][entry.entry_id], - description, - ) - for description in SENSOR_TYPES + GoalZeroSensor(entry.runtime_data, description) for description in SENSOR_TYPES ) diff --git a/homeassistant/components/goalzero/switch.py b/homeassistant/components/goalzero/switch.py index 9c0aee03b83..daff4ee5fec 100644 --- a/homeassistant/components/goalzero/switch.py +++ b/homeassistant/components/goalzero/switch.py @@ -5,11 +5,10 @@ from __future__ import annotations from typing import Any, cast from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .coordinator import GoalZeroConfigEntry from .entity import GoalZeroEntity SWITCH_TYPES: tuple[SwitchEntityDescription, ...] = ( @@ -29,15 +28,13 @@ SWITCH_TYPES: tuple[SwitchEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: GoalZeroConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Goal Zero Yeti switch.""" async_add_entities( - GoalZeroSwitch( - hass.data[DOMAIN][entry.entry_id], - description, - ) - for description in SWITCH_TYPES + GoalZeroSwitch(entry.runtime_data, description) for description in SWITCH_TYPES ) diff --git a/tests/components/goalzero/__init__.py b/tests/components/goalzero/__init__.py index d2e990ca122..30a7c92510e 100644 --- a/tests/components/goalzero/__init__.py +++ b/tests/components/goalzero/__init__.py @@ -3,8 +3,7 @@ from unittest.mock import AsyncMock, patch from homeassistant.components import dhcp -from homeassistant.components.goalzero import DOMAIN -from homeassistant.components.goalzero.const import DEFAULT_NAME +from homeassistant.components.goalzero.const import DEFAULT_NAME, DOMAIN from homeassistant.const import CONF_HOST, CONF_NAME from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import format_mac