Use runtime_data in velbus (#132988)

homeassistant/components/velbus/__init__.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import asyncio
+from dataclasses import dataclass
 import logging
 import os
 import shutil
@@ -34,6 +36,16 @@ PLATFORMS = [
 
 CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
 
+type VelbusConfigEntry = ConfigEntry[VelbusData]
+
+
+@dataclass
+class VelbusData:
+    """Runtime data for the Velbus config entry."""
+
+    controller: Velbus
+    connect_task: asyncio.Task
+
 
 async def velbus_connect_task(
     controller: Velbus, hass: HomeAssistant, entry_id: str
@@ -67,19 +79,14 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: VelbusConfigEntry) -> bool:
     """Establish connection with velbus."""
-    hass.data.setdefault(DOMAIN, {})
-
     controller = Velbus(
         entry.data[CONF_PORT],
         cache_dir=hass.config.path(STORAGE_DIR, f"velbuscache-{entry.entry_id}"),
     )
-    hass.data[DOMAIN][entry.entry_id] = {}
+    task = hass.async_create_task(velbus_connect_task(controller, hass, entry.entry_id))
-    hass.data[DOMAIN][entry.entry_id]["cntrl"] = controller
+    entry.runtime_data = VelbusData(controller=controller, connect_task=task)
-    hass.data[DOMAIN][entry.entry_id]["tsk"] = hass.async_create_task(
-        velbus_connect_task(controller, hass, entry.entry_id)
-    )
 
     _migrate_device_identifiers(hass, entry.entry_id)
 
@@ -88,17 +95,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: VelbusConfigEntry) -> bool:
     """Unload (close) the velbus connection."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    await hass.data[DOMAIN][entry.entry_id]["cntrl"].stop()
+    await entry.runtime_data.controller.stop()
-    hass.data[DOMAIN].pop(entry.entry_id)
-    if not hass.data[DOMAIN]:
-        hass.data.pop(DOMAIN)
     return unload_ok
 
 
-async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def async_remove_entry(hass: HomeAssistant, entry: VelbusConfigEntry) -> None:
     """Remove the velbus entry, so we also have to cleanup the cache dir."""
     await hass.async_add_executor_job(
         shutil.rmtree,
@@ -106,7 +110,9 @@ async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
     )
 
 
-async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_migrate_entry(
+    hass: HomeAssistant, config_entry: VelbusConfigEntry
+) -> bool:
     """Migrate old entry."""
     _LOGGER.debug("Migrating from version %s", config_entry.version)
     cache_path = hass.config.path(STORAGE_DIR, f"velbuscache-{config_entry.entry_id}/")

homeassistant/components/velbus/binary_sensor.py

@@ -3,24 +3,23 @@
 from velbusaio.channels import Button as VelbusButton
 
 from homeassistant.components.binary_sensor import BinarySensorEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import VelbusConfigEntry
 from .entity import VelbusEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: VelbusConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Velbus switch based on config_entry."""
-    await hass.data[DOMAIN][entry.entry_id]["tsk"]
+    await entry.runtime_data.connect_task
-    cntrl = hass.data[DOMAIN][entry.entry_id]["cntrl"]
     async_add_entities(
-        VelbusBinarySensor(channel) for channel in cntrl.get_all("binary_sensor")
+        VelbusBinarySensor(channel)
+        for channel in entry.runtime_data.controller.get_all_binary_sensor()
     )
 
 

homeassistant/components/velbus/button.py

@@ -8,24 +8,25 @@ from velbusaio.channels import (
 )
 
 from homeassistant.components.button import ButtonEntity
-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 . import VelbusConfigEntry
 from .entity import VelbusEntity, api_call
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: VelbusConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Velbus switch based on config_entry."""
-    await hass.data[DOMAIN][entry.entry_id]["tsk"]
+    await entry.runtime_data.connect_task
-    cntrl = hass.data[DOMAIN][entry.entry_id]["cntrl"]
+    async_add_entities(
-    async_add_entities(VelbusButton(channel) for channel in cntrl.get_all("button"))
+        VelbusButton(channel)
+        for channel in entry.runtime_data.controller.get_all_button()
+    )
 
 
 class VelbusButton(VelbusEntity, ButtonEntity):

homeassistant/components/velbus/climate.py

@@ -11,25 +11,27 @@ from homeassistant.components.climate import (
     ClimateEntityFeature,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import VelbusConfigEntry
 from .const import DOMAIN, PRESET_MODES
 from .entity import VelbusEntity, api_call
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: VelbusConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Velbus switch based on config_entry."""
-    await hass.data[DOMAIN][entry.entry_id]["tsk"]
+    await entry.runtime_data.connect_task
-    cntrl = hass.data[DOMAIN][entry.entry_id]["cntrl"]
+    async_add_entities(
-    async_add_entities(VelbusClimate(channel) for channel in cntrl.get_all("climate"))
+        VelbusClimate(channel)
+        for channel in entry.runtime_data.controller.get_all_climate()
+    )
 
 
 class VelbusClimate(VelbusEntity, ClimateEntity):

homeassistant/components/velbus/cover.py

@@ -11,23 +11,24 @@ from homeassistant.components.cover import (
     CoverEntity,
     CoverEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import VelbusConfigEntry
 from .entity import VelbusEntity, api_call
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: VelbusConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Velbus switch based on config_entry."""
-    await hass.data[DOMAIN][entry.entry_id]["tsk"]
+    await entry.runtime_data.connect_task
-    cntrl = hass.data[DOMAIN][entry.entry_id]["cntrl"]
+    async_add_entities(
-    async_add_entities(VelbusCover(channel) for channel in cntrl.get_all("cover"))
+        VelbusCover(channel)
+        for channel in entry.runtime_data.controller.get_all_cover()
+    )
 
 
 class VelbusCover(VelbusEntity, CoverEntity):

homeassistant/components/velbus/diagnostics.py

@@ -7,18 +7,17 @@ from typing import Any
 from velbusaio.channels import Channel as VelbusChannel
 from velbusaio.module import Module as VelbusModule
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntry
 
-from .const import DOMAIN
+from . import VelbusConfigEntry
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: VelbusConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    controller = hass.data[DOMAIN][entry.entry_id]["cntrl"]
+    controller = entry.runtime_data.controller
     data: dict[str, Any] = {"entry": entry.as_dict(), "modules": []}
     for module in controller.get_modules().values():
         data["modules"].append(_build_module_diagnostics_info(module))
@@ -26,10 +25,10 @@ async def async_get_config_entry_diagnostics(
 
 
 async def async_get_device_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry, device: DeviceEntry
+    hass: HomeAssistant, entry: VelbusConfigEntry, device: DeviceEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a device entry."""
-    controller = hass.data[DOMAIN][entry.entry_id]["cntrl"]
+    controller = entry.runtime_data.controller
     channel = list(next(iter(device.identifiers)))[1]
     modules = controller.get_modules()
     return _build_module_diagnostics_info(modules[int(channel)])

homeassistant/components/velbus/light.py

@@ -20,28 +20,30 @@ from homeassistant.components.light import (
     LightEntity,
     LightEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import VelbusConfigEntry
 from .entity import VelbusEntity, api_call
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: VelbusConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Velbus switch based on config_entry."""
-    await hass.data[DOMAIN][entry.entry_id]["tsk"]
+    await entry.runtime_data.connect_task
-    cntrl = hass.data[DOMAIN][entry.entry_id]["cntrl"]
     entities: list[Entity] = [
-        VelbusLight(channel) for channel in cntrl.get_all("light")
+        VelbusLight(channel)
+        for channel in entry.runtime_data.controller.get_all_light()
     ]
-    entities.extend(VelbusButtonLight(channel) for channel in cntrl.get_all("led"))
+    entities.extend(
+        VelbusButtonLight(channel)
+        for channel in entry.runtime_data.controller.get_all_led()
+    )
     async_add_entities(entities)
 
 

homeassistant/components/velbus/quality_scale.yaml

@@ -23,7 +23,7 @@ rules:
   entity-event-setup: todo
   entity-unique-id: done
   has-entity-name: todo
-  runtime-data: todo
+  runtime-data: done
   test-before-configure: done
   test-before-setup: todo
   unique-config-entry:

homeassistant/components/velbus/select.py

@@ -3,24 +3,25 @@
 from velbusaio.channels import SelectedProgram
 
 from homeassistant.components.select import SelectEntity
-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 . import VelbusConfigEntry
 from .entity import VelbusEntity, api_call
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: VelbusConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Velbus select based on config_entry."""
-    await hass.data[DOMAIN][entry.entry_id]["tsk"]
+    await entry.runtime_data.connect_task
-    cntrl = hass.data[DOMAIN][entry.entry_id]["cntrl"]
+    async_add_entities(
-    async_add_entities(VelbusSelect(channel) for channel in cntrl.get_all("select"))
+        VelbusSelect(channel)
+        for channel in entry.runtime_data.controller.get_all_select()
+    )
 
 
 class VelbusSelect(VelbusEntity, SelectEntity):

homeassistant/components/velbus/sensor.py

@@ -9,24 +9,22 @@ from homeassistant.components.sensor import (
     SensorEntity,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import VelbusConfigEntry
 from .entity import VelbusEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: VelbusConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Velbus switch based on config_entry."""
-    await hass.data[DOMAIN][entry.entry_id]["tsk"]
+    await entry.runtime_data.connect_task
-    cntrl = hass.data[DOMAIN][entry.entry_id]["cntrl"]
     entities = []
-    for channel in cntrl.get_all("sensor"):
+    for channel in entry.runtime_data.controller.get_all_sensor():
         entities.append(VelbusSensor(channel))
         if channel.is_counter_channel():
             entities.append(VelbusSensor(channel, True))

homeassistant/components/velbus/services.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 from contextlib import suppress
 import os
 import shutil
+from typing import TYPE_CHECKING
 
 import voluptuous as vol
 
@@ -13,6 +14,9 @@ from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.storage import STORAGE_DIR
 
+if TYPE_CHECKING:
+    from . import VelbusConfigEntry
+
 from .const import (
     CONF_INTERFACE,
     CONF_MEMO_TEXT,
@@ -35,20 +39,32 @@ def setup_services(hass: HomeAssistant) -> None:
             "The interface provided is not defined as a port in a Velbus integration"
         )
 
+    def get_config_entry(interface: str) -> VelbusConfigEntry | None:
+        for config_entry in hass.config_entries.async_entries(DOMAIN):
+            if "port" in config_entry.data and config_entry.data["port"] == interface:
+                return config_entry
+        return None
+
     async def scan(call: ServiceCall) -> None:
-        await hass.data[DOMAIN][call.data[CONF_INTERFACE]]["cntrl"].scan()
+        """Handle a scan service call."""
+        entry = get_config_entry(call.data[CONF_INTERFACE])
+        if entry:
+            await entry.runtime_data.controller.scan()
 
     async def syn_clock(call: ServiceCall) -> None:
-        await hass.data[DOMAIN][call.data[CONF_INTERFACE]]["cntrl"].sync_clock()
+        """Handle a sync clock service call."""
+        entry = get_config_entry(call.data[CONF_INTERFACE])
+        if entry:
+            await entry.runtime_data.controller.sync_clock()
 
     async def set_memo_text(call: ServiceCall) -> None:
         """Handle Memo Text service call."""
-        memo_text = call.data[CONF_MEMO_TEXT]
+        entry = get_config_entry(call.data[CONF_INTERFACE])
-        await (
+        if entry:
-            hass.data[DOMAIN][call.data[CONF_INTERFACE]]["cntrl"]
+            memo_text = call.data[CONF_MEMO_TEXT]
-            .get_module(call.data[CONF_ADDRESS])
+            module = entry.runtime_data.controller.get_module(call.data[CONF_ADDRESS])
-            .set_memo_text(memo_text.async_render())
+            if module:
-        )
+                await module.set_memo_text(memo_text.async_render())
 
     async def clear_cache(call: ServiceCall) -> None:
         """Handle a clear cache service call."""

homeassistant/components/velbus/switch.py

@@ -5,23 +5,24 @@ from typing import Any
 from velbusaio.channels import Relay as VelbusRelay
 
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import VelbusConfigEntry
 from .entity import VelbusEntity, api_call
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: VelbusConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Velbus switch based on config_entry."""
-    await hass.data[DOMAIN][entry.entry_id]["tsk"]
+    await entry.runtime_data.connect_task
-    cntrl = hass.data[DOMAIN][entry.entry_id]["cntrl"]
+    async_add_entities(
-    async_add_entities(VelbusSwitch(channel) for channel in cntrl.get_all("switch"))
+        VelbusSwitch(channel)
+        for channel in entry.runtime_data.controller.get_all_switch()
+    )
 
 
 class VelbusSwitch(VelbusEntity, SwitchEntity):