Migrate Z-Wave.Me to use runtime_data (#168562)

Author: mib1185

homeassistant/components/zwave_me/__init__.py

@@ -1,103 +1,37 @@
 """The Z-Wave-Me WS integration."""
 
-from zwave_me_ws import ZWaveMe, ZWaveMeData
-
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_TOKEN, CONF_URL
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import device_registry as dr
-from homeassistant.helpers.dispatcher import dispatcher_send
-
-from .const import DOMAIN, PLATFORMS, ZWaveMePlatform
 
-ZWAVE_ME_PLATFORMS = [platform.value for platform in ZWaveMePlatform]
+from .const import PLATFORMS
+from .controller import ZWaveMeConfigEntry, ZWaveMeController
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: ZWaveMeConfigEntry) -> bool:
     """Set up Z-Wave-Me from a config entry."""
-    hass.data.setdefault(DOMAIN, {})
-    controller = hass.data[DOMAIN][entry.entry_id] = ZWaveMeController(hass, entry)
-    if await controller.async_establish_connection():
-        await async_setup_platforms(hass, entry, controller)
-        registry = dr.async_get(hass)
-        controller.remove_stale_devices(registry)
-        return True
-    raise ConfigEntryNotReady
+    controller = ZWaveMeController(hass, entry)
+
+    if not await controller.async_establish_connection():
+        raise ConfigEntryNotReady
+
+    entry.runtime_data = controller
+    await async_setup_platforms(hass, entry, controller)
+    registry = dr.async_get(hass)
+    controller.remove_stale_devices(registry)
+    return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: ZWaveMeConfigEntry) -> bool:
     """Unload a config entry."""
 
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
-        controller = hass.data[DOMAIN].pop(entry.entry_id)
-        await controller.zwave_api.close_ws()
+        await entry.runtime_data.zwave_api.close_ws()
     return unload_ok
 
 
-class ZWaveMeController:
-    """Main ZWave-Me API class."""
-
-    def __init__(self, hass: HomeAssistant, config: ConfigEntry) -> None:
-        """Create the API instance."""
-        self.device_ids: set = set()
-        self._hass = hass
-        self.config = config
-        self.zwave_api = ZWaveMe(
-            on_device_create=self.on_device_create,
-            on_device_update=self.on_device_update,
-            on_device_remove=self.on_device_unavailable,
-            on_device_destroy=self.on_device_destroy,
-            on_new_device=self.add_device,
-            token=self.config.data[CONF_TOKEN],
-            url=self.config.data[CONF_URL],
-            platforms=ZWAVE_ME_PLATFORMS,
-        )
-        self.platforms_inited = False
-
-    async def async_establish_connection(self):
-        """Get connection status."""
-        return await self.zwave_api.get_connection()
-
-    def add_device(self, device: ZWaveMeData) -> None:
-        """Send signal to create device."""
-        if device.id in self.device_ids:
-            dispatcher_send(self._hass, f"ZWAVE_ME_INFO_{device.id}", device)
-        else:
-            dispatcher_send(
-                self._hass, f"ZWAVE_ME_NEW_{device.deviceType.upper()}", device
-            )
-            self.device_ids.add(device.id)
-
-    def on_device_create(self, devices: list[ZWaveMeData]) -> None:
-        """Create multiple devices."""
-        for device in devices:
-            if device.deviceType in ZWAVE_ME_PLATFORMS and self.platforms_inited:
-                self.add_device(device)
-
-    def on_device_update(self, new_info: ZWaveMeData) -> None:
-        """Send signal to update device."""
-        dispatcher_send(self._hass, f"ZWAVE_ME_INFO_{new_info.id}", new_info)
-
-    def on_device_unavailable(self, device_id: str) -> None:
-        """Send signal to set device unavailable."""
-        dispatcher_send(self._hass, f"ZWAVE_ME_UNAVAILABLE_{device_id}")
-
-    def on_device_destroy(self, device_id: str) -> None:
-        """Send signal to destroy device."""
-        dispatcher_send(self._hass, f"ZWAVE_ME_DESTROY_{device_id}")
-
-    def remove_stale_devices(self, registry: dr.DeviceRegistry):
-        """Remove old-format devices in the registry."""
-        for device_id in self.device_ids:
-            device = registry.async_get_device(
-                identifiers={(DOMAIN, f"{self.config.unique_id}-{device_id}")}
-            )
-            if device is not None:
-                registry.async_remove_device(device.id)
-
-
 async def async_setup_platforms(
     hass: HomeAssistant, entry: ConfigEntry, controller: ZWaveMeController
 ) -> None:

homeassistant/components/zwave_me/binary_sensor.py

@@ -9,13 +9,12 @@
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import ZWaveMeController
-from .const import DOMAIN, ZWaveMePlatform
+from .const import ZWaveMePlatform
+from .controller import ZWaveMeConfigEntry, ZWaveMeController
 from .entity import ZWaveMeEntity
 
 BINARY_SENSORS_MAP: dict[str, BinarySensorEntityDescription] = {
@@ -32,22 +31,22 @@
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: ZWaveMeConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the binary sensor platform."""
 
     @callback
     def add_new_device(new_device: ZWaveMeData) -> None:
-        controller: ZWaveMeController = hass.data[DOMAIN][config_entry.entry_id]
-        description = BINARY_SENSORS_MAP.get(
-            new_device.probeType, BINARY_SENSORS_MAP["generic"]
-        )
-        sensor = ZWaveMeBinarySensor(controller, new_device, description)
-
         async_add_entities(
             [
-                sensor,
+                ZWaveMeBinarySensor(
+                    config_entry.runtime_data,
+                    new_device,
+                    BINARY_SENSORS_MAP.get(
+                        new_device.probeType, BINARY_SENSORS_MAP["generic"]
+                    ),
+                )
             ]
         )
 

homeassistant/components/zwave_me/button.py

@@ -1,34 +1,27 @@
 """Representation of a toggleButton."""
 
 from homeassistant.components.button import ButtonEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import DOMAIN, ZWaveMePlatform
+from .const import ZWaveMePlatform
+from .controller import ZWaveMeConfigEntry
 from .entity import ZWaveMeEntity
 
 DEVICE_NAME = ZWaveMePlatform.BUTTON
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: ZWaveMeConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the number platform."""
 
     @callback
     def add_new_device(new_device):
-        controller = hass.data[DOMAIN][config_entry.entry_id]
-        button = ZWaveMeButton(controller, new_device)
-
-        async_add_entities(
-            [
-                button,
-            ]
-        )
+        async_add_entities([ZWaveMeButton(config_entry.runtime_data, new_device)])
 
     config_entry.async_on_unload(
         async_dispatcher_connect(

homeassistant/components/zwave_me/climate.py

@@ -11,13 +11,13 @@
     ClimateEntityFeature,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import DOMAIN, ZWaveMePlatform
+from .const import ZWaveMePlatform
+from .controller import ZWaveMeConfigEntry
 from .entity import ZWaveMeEntity
 
 TEMPERATURE_DEFAULT_STEP = 0.5
@@ -27,22 +27,15 @@
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: ZWaveMeConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the climate platform."""
 
     @callback
     def add_new_device(new_device: ZWaveMeData) -> None:
         """Add a new device."""
-        controller = hass.data[DOMAIN][config_entry.entry_id]
-        climate = ZWaveMeClimate(controller, new_device)
-
-        async_add_entities(
-            [
-                climate,
-            ]
-        )
+        async_add_entities([ZWaveMeClimate(config_entry.runtime_data, new_device)])
 
     config_entry.async_on_unload(
         async_dispatcher_connect(

homeassistant/components/zwave_me/controller.py

@@ -0,0 +1,77 @@
+"""The Z-Wave-Me WS controller."""
+
+from zwave_me_ws import ZWaveMe, ZWaveMeData
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_TOKEN, CONF_URL
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers.dispatcher import dispatcher_send
+
+from .const import DOMAIN, ZWaveMePlatform
+
+type ZWaveMeConfigEntry = ConfigEntry[ZWaveMeController]
+
+ZWAVE_ME_PLATFORMS = [platform.value for platform in ZWaveMePlatform]
+
+
+class ZWaveMeController:
+    """Main ZWave-Me API class."""
+
+    def __init__(self, hass: HomeAssistant, config: ZWaveMeConfigEntry) -> None:
+        """Create the API instance."""
+        self.device_ids: set[str] = set()
+        self._hass = hass
+        self.config = config
+        self.zwave_api = ZWaveMe(
+            on_device_create=self.on_device_create,
+            on_device_update=self.on_device_update,
+            on_device_remove=self.on_device_unavailable,
+            on_device_destroy=self.on_device_destroy,
+            on_new_device=self.add_device,
+            token=self.config.data[CONF_TOKEN],
+            url=self.config.data[CONF_URL],
+            platforms=ZWAVE_ME_PLATFORMS,
+        )
+        self.platforms_inited = False
+
+    async def async_establish_connection(self) -> bool:
+        """Get connection status."""
+        return await self.zwave_api.get_connection()
+
+    def add_device(self, device: ZWaveMeData) -> None:
+        """Send signal to create device."""
+        if device.id in self.device_ids:
+            dispatcher_send(self._hass, f"ZWAVE_ME_INFO_{device.id}", device)
+        else:
+            dispatcher_send(
+                self._hass, f"ZWAVE_ME_NEW_{device.deviceType.upper()}", device
+            )
+            self.device_ids.add(device.id)
+
+    def on_device_create(self, devices: list[ZWaveMeData]) -> None:
+        """Create multiple devices."""
+        for device in devices:
+            if device.deviceType in ZWAVE_ME_PLATFORMS and self.platforms_inited:
+                self.add_device(device)
+
+    def on_device_update(self, new_info: ZWaveMeData) -> None:
+        """Send signal to update device."""
+        dispatcher_send(self._hass, f"ZWAVE_ME_INFO_{new_info.id}", new_info)
+
+    def on_device_unavailable(self, device_id: str) -> None:
+        """Send signal to set device unavailable."""
+        dispatcher_send(self._hass, f"ZWAVE_ME_UNAVAILABLE_{device_id}")
+
+    def on_device_destroy(self, device_id: str) -> None:
+        """Send signal to destroy device."""
+        dispatcher_send(self._hass, f"ZWAVE_ME_DESTROY_{device_id}")
+
+    def remove_stale_devices(self, registry: dr.DeviceRegistry):
+        """Remove old-format devices in the registry."""
+        for device_id in self.device_ids:
+            device = registry.async_get_device(
+                identifiers={(DOMAIN, f"{self.config.unique_id}-{device_id}")}
+            )
+            if device is not None:
+                registry.async_remove_device(device.id)

homeassistant/components/zwave_me/cover.py

@@ -9,34 +9,27 @@
     CoverEntity,
     CoverEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import DOMAIN, ZWaveMePlatform
+from .const import ZWaveMePlatform
+from .controller import ZWaveMeConfigEntry
 from .entity import ZWaveMeEntity
 
 DEVICE_NAME = ZWaveMePlatform.COVER
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: ZWaveMeConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the cover platform."""
 
     @callback
     def add_new_device(new_device):
-        controller = hass.data[DOMAIN][config_entry.entry_id]
-        cover = ZWaveMeCover(controller, new_device)
-
-        async_add_entities(
-            [
-                cover,
-            ]
-        )
+        async_add_entities([ZWaveMeCover(config_entry.runtime_data, new_device)])
 
     config_entry.async_on_unload(
         async_dispatcher_connect(

homeassistant/components/zwave_me/entity.py

@@ -8,12 +8,13 @@
 from homeassistant.helpers.entity import Entity
 
 from .const import DOMAIN
+from .controller import ZWaveMeController
 
 
 class ZWaveMeEntity(Entity):
     """Representation of a ZWaveMe device."""
 
-    def __init__(self, controller, device):
+    def __init__(self, controller: ZWaveMeController, device: ZWaveMeData) -> None:
         """Initialize the device."""
         self.controller = controller
         self.device = device