Use runtime_data in crownstone (#136406)

* Use runtime_data in crownstone

* Move some logic into __init__

* Remove underscore in async_update_listener

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)

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)

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)

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] = []
 

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))
         ),