Register Wemo fan services with entity service helper

homeassistant/components/wemo/__init__.py

@@ -8,8 +8,10 @@
 
 from homeassistant import config_entries
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.singleton import singleton
 
 from .const import DOMAIN
 
@@ -68,11 +70,16 @@ def coerce_host_port(value):
 )
 
 
+@singleton("f{__name__}.registry")
+async def async_get_registry(hass: HomeAssistant) -> pywemo.SubscriptionRegistry:
+    """Return the pywemo subscription registry."""
+    return await hass.async_add_executor_job(pywemo.SubscriptionRegistry)
+
+
 async def async_setup(hass, config):
     """Set up for WeMo devices."""
     hass.data[DOMAIN] = {
         "config": config.get(DOMAIN, {}),
-        "registry": None,
         "pending": {},
     }
 
@@ -91,7 +98,7 @@ async def async_setup_entry(hass, entry):
     config = hass.data[DOMAIN].pop("config")
 
     # Keep track of WeMo device subscriptions for push updates
-    registry = hass.data[DOMAIN]["registry"] = pywemo.SubscriptionRegistry()
+    registry = await async_get_registry(hass)
     await hass.async_add_executor_job(registry.start)
 
     def stop_wemo(event):

homeassistant/components/wemo/entity.py

@@ -8,6 +8,7 @@
 
 from homeassistant.helpers.entity import Entity
 
+from . import async_get_registry
 from .const import DOMAIN as WEMO_DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
@@ -97,13 +98,13 @@ async def async_added_to_hass(self) -> None:
         """Wemo device added to Home Assistant."""
         await super().async_added_to_hass()
 
-        registry = self.hass.data[WEMO_DOMAIN]["registry"]
+        registry = await async_get_registry(self.hass)
         await self.hass.async_add_executor_job(registry.register, self.wemo)
         registry.on(self.wemo, None, self._subscription_callback)
 
     async def async_will_remove_from_hass(self) -> None:
         """Wemo device removed from hass."""
-        registry = self.hass.data[WEMO_DOMAIN]["registry"]
+        registry = await async_get_registry(self.hass)
         await self.hass.async_add_executor_job(registry.unregister, self.wemo)
 
     def _subscription_callback(

homeassistant/components/wemo/fan.py

@@ -2,6 +2,7 @@
 import asyncio
 from datetime import timedelta
 import logging
+from typing import List
 
 from pywemo.ouimeaux_device.api.service import ActionException
 import voluptuous as vol
@@ -15,8 +16,10 @@
     FanEntity,
 )
 from homeassistant.const import ATTR_ENTITY_ID
+from homeassistant.core import HomeAssistant
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.singleton import singleton
 
 from .const import (
     DOMAIN as WEMO_DOMAIN,
@@ -93,24 +96,14 @@
 RESET_FILTER_LIFE_SCHEMA = vol.Schema({vol.Required(ATTR_ENTITY_ID): cv.entity_ids})
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up WeMo binary sensors."""
-    entities = []
+@singleton(f"{__name__}.service")
+async def _async_get_humidifiers(hass: HomeAssistant) -> List["WemoHumidifier"]:
+    """Singleton containing humidifier services.
 
-    async def _discovered_wemo(device):
-        """Handle a discovered Wemo device."""
-        entity = WemoHumidifier(device)
-        entities.append(entity)
-        async_add_entities([entity])
-
-    async_dispatcher_connect(hass, f"{WEMO_DOMAIN}.fan", _discovered_wemo)
-
-    await asyncio.gather(
-        *[
-            _discovered_wemo(device)
-            for device in hass.data[WEMO_DOMAIN]["pending"].pop("fan")
-        ]
-    )
+    Returns:
+      List of humidifiers used by the services.
+    """
+    entities = []
 
     def service_handle(service):
         """Handle the WeMo humidifier services."""
@@ -142,6 +135,25 @@ def service_handle(service):
         schema=RESET_FILTER_LIFE_SCHEMA,
     )
 
+    return entities
+
+
+async def async_setup_entry(hass, config_entry, async_add_entities):
+    """Set up WeMo binary sensors."""
+
+    async def _discovered_wemo(device):
+        """Handle a discovered Wemo device."""
+        async_add_entities([WemoHumidifier(device)])
+
+    async_dispatcher_connect(hass, f"{WEMO_DOMAIN}.fan", _discovered_wemo)
+
+    await asyncio.gather(
+        *[
+            _discovered_wemo(device)
+            for device in hass.data[WEMO_DOMAIN]["pending"].pop("fan")
+        ]
+    )
+
 
 class WemoHumidifier(WemoSubscriptionEntity, FanEntity):
     """Representation of a WeMo humidifier."""
@@ -189,6 +201,16 @@ def supported_features(self) -> int:
         """Flag supported features."""
         return SUPPORTED_FEATURES
 
+    async def async_added_to_hass(self) -> None:
+        """Wemo humidifier added to Home Assistant."""
+        await super().async_added_to_hass()
+        (await _async_get_humidifiers(self.hass)).append(self)
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Wemo humidifier removed from hass."""
+        await super().async_will_remove_from_hass()
+        (await _async_get_humidifiers(self.hass)).remove(self)
+
     def _update(self, force_update=True):
         """Update the device state."""
         try: