From ffdb686363d61ac82cd2daf68ea81b18f1e763e8 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Mon, 27 Jan 2025 10:15:53 +0100 Subject: [PATCH] Use runtime_data in crownstone (#136406) * Use runtime_data in crownstone * Move some logic into __init__ * Remove underscore in async_update_listener --- .../components/crownstone/__init__.py | 39 +++++++++++++----- .../components/crownstone/config_flow.py | 12 +++--- .../components/crownstone/entry_manager.py | 40 +++++-------------- homeassistant/components/crownstone/light.py | 12 ++---- .../components/crownstone/test_config_flow.py | 15 +++---- 5 files changed, 55 insertions(+), 63 deletions(-) diff --git a/homeassistant/components/crownstone/__init__.py b/homeassistant/components/crownstone/__init__.py index e1443eb9516..8f5739f9172 100644 --- a/homeassistant/components/crownstone/__init__.py +++ b/homeassistant/components/crownstone/__init__.py @@ -2,25 +2,42 @@ from __future__ import annotations -from homeassistant.config_entries import ConfigEntry +from homeassistant.const import EVENT_HOMEASSISTANT_STOP from homeassistant.core import HomeAssistant -from .const import DOMAIN -from .entry_manager import CrownstoneEntryManager +from .const import PLATFORMS +from .entry_manager import CrownstoneConfigEntry, CrownstoneEntryManager -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: CrownstoneConfigEntry) -> bool: """Initiate setup for a Crownstone config entry.""" manager = CrownstoneEntryManager(hass, entry) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = manager + if not await manager.async_setup(): + return False - return await manager.async_setup() + entry.runtime_data = manager + + await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) + + # HA specific listeners + entry.async_on_unload(entry.add_update_listener(async_update_listener)) + entry.async_on_unload( + hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, manager.on_shutdown) + ) + + return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: CrownstoneConfigEntry) -> bool: """Unload a config entry.""" - unload_ok: bool = await hass.data[DOMAIN][entry.entry_id].async_unload() - if len(hass.data[DOMAIN]) == 0: - hass.data.pop(DOMAIN) - return unload_ok + entry.runtime_data.async_unload() + + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) + + +async def async_update_listener( + hass: HomeAssistant, entry: CrownstoneConfigEntry +) -> None: + """Handle options update.""" + await hass.config_entries.async_reload(entry.entry_id) diff --git a/homeassistant/components/crownstone/config_flow.py b/homeassistant/components/crownstone/config_flow.py index 2a96098421a..5f5af4f51a4 100644 --- a/homeassistant/components/crownstone/config_flow.py +++ b/homeassistant/components/crownstone/config_flow.py @@ -16,7 +16,6 @@ import voluptuous as vol from homeassistant.components import usb from homeassistant.config_entries import ( - ConfigEntry, ConfigEntryBaseFlow, ConfigFlow, ConfigFlowResult, @@ -37,6 +36,7 @@ from .const import ( MANUAL_PATH, REFRESH_LIST, ) +from .entry_manager import CrownstoneConfigEntry from .helpers import list_ports_as_str CONFIG_FLOW = "config_flow" @@ -140,7 +140,7 @@ class CrownstoneConfigFlowHandler(BaseCrownstoneFlowHandler, ConfigFlow, domain= @staticmethod @callback def async_get_options_flow( - config_entry: ConfigEntry, + config_entry: CrownstoneConfigEntry, ) -> CrownstoneOptionsFlowHandler: """Return the Crownstone options.""" return CrownstoneOptionsFlowHandler(config_entry) @@ -210,7 +210,9 @@ class CrownstoneConfigFlowHandler(BaseCrownstoneFlowHandler, ConfigFlow, domain= class CrownstoneOptionsFlowHandler(BaseCrownstoneFlowHandler, OptionsFlow): """Handle Crownstone options.""" - def __init__(self, config_entry: ConfigEntry) -> None: + config_entry: CrownstoneConfigEntry + + def __init__(self, config_entry: CrownstoneConfigEntry) -> None: """Initialize Crownstone options.""" super().__init__(OPTIONS_FLOW, self.async_create_new_entry) self.options = config_entry.options.copy() @@ -219,9 +221,7 @@ class CrownstoneOptionsFlowHandler(BaseCrownstoneFlowHandler, OptionsFlow): self, user_input: dict[str, Any] | None = None ) -> ConfigFlowResult: """Manage Crownstone options.""" - self.cloud: CrownstoneCloud = self.hass.data[DOMAIN][ - self.config_entry.entry_id - ].cloud + self.cloud = self.config_entry.runtime_data.cloud spheres = {sphere.name: sphere.cloud_id for sphere in self.cloud.cloud_data} usb_path = self.config_entry.options.get(CONF_USB_PATH) diff --git a/homeassistant/components/crownstone/entry_manager.py b/homeassistant/components/crownstone/entry_manager.py index efee05a19c8..e414e3c7055 100644 --- a/homeassistant/components/crownstone/entry_manager.py +++ b/homeassistant/components/crownstone/entry_manager.py @@ -16,7 +16,7 @@ from crownstone_uart.Exceptions import UartException from homeassistant.components import persistent_notification from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, EVENT_HOMEASSISTANT_STOP +from homeassistant.const import CONF_EMAIL, CONF_PASSWORD from homeassistant.core import Event, HomeAssistant, callback from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers import aiohttp_client @@ -26,7 +26,6 @@ from .const import ( CONF_USB_PATH, CONF_USB_SPHERE, DOMAIN, - PLATFORMS, PROJECT_NAME, SSE_LISTENERS, UART_LISTENERS, @@ -36,6 +35,8 @@ from .listeners import setup_sse_listeners, setup_uart_listeners _LOGGER = logging.getLogger(__name__) +type CrownstoneConfigEntry = ConfigEntry[CrownstoneEntryManager] + class CrownstoneEntryManager: """Manage a Crownstone config entry.""" @@ -44,7 +45,9 @@ class CrownstoneEntryManager: cloud: CrownstoneCloud sse: CrownstoneSSEAsync - def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None: + def __init__( + self, hass: HomeAssistant, config_entry: CrownstoneConfigEntry + ) -> None: """Initialize the hub.""" self.hass = hass self.config_entry = config_entry @@ -100,18 +103,6 @@ class CrownstoneEntryManager: # Makes HA aware of the Crownstone environment HA is placed in, a user can have multiple self.usb_sphere_id = self.config_entry.options[CONF_USB_SPHERE] - await self.hass.config_entries.async_forward_entry_setups( - self.config_entry, PLATFORMS - ) - - # HA specific listeners - self.config_entry.async_on_unload( - self.config_entry.add_update_listener(_async_update_listener) - ) - self.config_entry.async_on_unload( - self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, self.on_shutdown) - ) - return True async def async_process_events(self, sse_client: CrownstoneSSEAsync) -> None: @@ -161,11 +152,12 @@ class CrownstoneEntryManager: setup_uart_listeners(self) - async def async_unload(self) -> bool: + @callback + def async_unload(self) -> None: """Unload the current config entry.""" # Authentication failed if self.cloud.cloud_data is None: - return True + return self.sse.close_client() for sse_unsub in self.listeners[SSE_LISTENERS]: @@ -176,23 +168,9 @@ class CrownstoneEntryManager: for subscription_id in self.listeners[UART_LISTENERS]: UartEventBus.unsubscribe(subscription_id) - unload_ok = await self.hass.config_entries.async_unload_platforms( - self.config_entry, PLATFORMS - ) - - if unload_ok: - self.hass.data[DOMAIN].pop(self.config_entry.entry_id) - - return unload_ok - @callback def on_shutdown(self, _: Event) -> None: """Close all IO connections.""" self.sse.close_client() if self.uart: self.uart.stop() - - -async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None: - """Handle options update.""" - await hass.config_entries.async_reload(entry.entry_id) diff --git a/homeassistant/components/crownstone/light.py b/homeassistant/components/crownstone/light.py index 16faa3a36d2..70b7631fe6b 100644 --- a/homeassistant/components/crownstone/light.py +++ b/homeassistant/components/crownstone/light.py @@ -3,7 +3,7 @@ from __future__ import annotations from functools import partial -from typing import TYPE_CHECKING, Any +from typing import Any from crownstone_cloud.cloud_models.crownstones import Crownstone from crownstone_cloud.const import DIMMING_ABILITY @@ -11,7 +11,6 @@ from crownstone_cloud.exceptions import CrownstoneAbilityError from crownstone_uart import CrownstoneUart from homeassistant.components.light import ATTR_BRIGHTNESS, ColorMode, LightEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.dispatcher import async_dispatcher_connect @@ -20,24 +19,21 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from .const import ( CROWNSTONE_INCLUDE_TYPES, CROWNSTONE_SUFFIX, - DOMAIN, SIG_CROWNSTONE_STATE_UPDATE, SIG_UART_STATE_CHANGE, ) from .entity import CrownstoneEntity +from .entry_manager import CrownstoneConfigEntry from .helpers import map_from_to -if TYPE_CHECKING: - from .entry_manager import CrownstoneEntryManager - async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: CrownstoneConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up crownstones from a config entry.""" - manager: CrownstoneEntryManager = hass.data[DOMAIN][config_entry.entry_id] + manager = config_entry.runtime_data entities: list[CrownstoneLightEntity] = [] diff --git a/tests/components/crownstone/test_config_flow.py b/tests/components/crownstone/test_config_flow.py index a38a04cb2ad..c3bb17cb6d6 100644 --- a/tests/components/crownstone/test_config_flow.py +++ b/tests/components/crownstone/test_config_flow.py @@ -163,7 +163,7 @@ async def start_config_flow(hass: HomeAssistant, mocked_cloud: MagicMock): async def start_options_flow( - hass: HomeAssistant, entry_id: str, mocked_manager: MagicMock + hass: HomeAssistant, entry: MockConfigEntry, mocked_manager: MagicMock ): """Patch CrownstoneEntryManager and start the flow.""" # set up integration @@ -171,9 +171,10 @@ async def start_options_flow( "homeassistant.components.crownstone.CrownstoneEntryManager", return_value=mocked_manager, ): - await hass.config_entries.async_setup(entry_id) + await hass.config_entries.async_setup(entry.entry_id) - return await hass.config_entries.options.async_init(entry_id) + entry.runtime_data = mocked_manager + return await hass.config_entries.options.async_init(entry.entry_id) async def test_no_user_input( @@ -413,7 +414,7 @@ async def test_options_flow_setup_usb( result = await start_options_flow( hass, - entry.entry_id, + entry, get_mocked_crownstone_entry_manager( get_mocked_crownstone_cloud(create_mocked_spheres(2)) ), @@ -490,7 +491,7 @@ async def test_options_flow_remove_usb(hass: HomeAssistant) -> None: result = await start_options_flow( hass, - entry.entry_id, + entry, get_mocked_crownstone_entry_manager( get_mocked_crownstone_cloud(create_mocked_spheres(2)) ), @@ -543,7 +544,7 @@ async def test_options_flow_manual_usb_path( result = await start_options_flow( hass, - entry.entry_id, + entry, get_mocked_crownstone_entry_manager( get_mocked_crownstone_cloud(create_mocked_spheres(1)) ), @@ -602,7 +603,7 @@ async def test_options_flow_change_usb_sphere(hass: HomeAssistant) -> None: result = await start_options_flow( hass, - entry.entry_id, + entry, get_mocked_crownstone_entry_manager( get_mocked_crownstone_cloud(create_mocked_spheres(3)) ),