Use asyncio in XiaomiAqara instead of threading (#74979)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/xiaomi_aqara/__init__.py

@@ -1,12 +1,13 @@
 """Support for Xiaomi Gateways."""
+import asyncio
 from datetime import timedelta
 import logging
 
 import voluptuous as vol
-from xiaomi_gateway import XiaomiGateway, XiaomiGatewayDiscovery
+from xiaomi_gateway import AsyncXiaomiGatewayMulticast, XiaomiGateway
 
 from homeassistant.components import persistent_notification
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import (
     ATTR_BATTERY_LEVEL,
     ATTR_DEVICE_ID,
@@ -34,6 +35,8 @@ from .const import (
     DEFAULT_DISCOVERY_RETRY,
     DOMAIN,
     GATEWAYS_KEY,
+    KEY_SETUP_LOCK,
+    KEY_UNSUB_STOP,
     LISTENER_KEY,
 )
 
@@ -143,6 +146,7 @@ def setup(hass: HomeAssistant, config: ConfigType) -> bool:
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up the xiaomi aqara components from a config entry."""
     hass.data.setdefault(DOMAIN, {})
+    setup_lock = hass.data[DOMAIN].setdefault(KEY_SETUP_LOCK, asyncio.Lock())
     hass.data[DOMAIN].setdefault(GATEWAYS_KEY, {})
 
     # Connect to Xiaomi Aqara Gateway
@@ -158,24 +162,28 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
     hass.data[DOMAIN][GATEWAYS_KEY][entry.entry_id] = xiaomi_gateway
 
-    gateway_discovery = hass.data[DOMAIN].setdefault(
+    async with setup_lock:
-        LISTENER_KEY,
+        if LISTENER_KEY not in hass.data[DOMAIN]:
-        XiaomiGatewayDiscovery(hass.add_job, [], entry.data[CONF_INTERFACE]),
+            multicast = AsyncXiaomiGatewayMulticast(
+                interface=entry.data[CONF_INTERFACE]
             )
+            hass.data[DOMAIN][LISTENER_KEY] = multicast
 
-    if len(hass.data[DOMAIN][GATEWAYS_KEY]) == 1:
             # start listining for local pushes (only once)
-        await hass.async_add_executor_job(gateway_discovery.listen)
+            await multicast.start_listen()
 
             # register stop callback to shutdown listining for local pushes
+            @callback
             def stop_xiaomi(event):
                 """Stop Xiaomi Socket."""
                 _LOGGER.debug("Shutting down Xiaomi Gateway Listener")
-            gateway_discovery.stop_listen()
+                multicast.stop_listen()
 
-        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, stop_xiaomi)
+            unsub = hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, stop_xiaomi)
+            hass.data[DOMAIN][KEY_UNSUB_STOP] = unsub
 
-    gateway_discovery.gateways[entry.data[CONF_HOST]] = xiaomi_gateway
+    multicast = hass.data[DOMAIN][LISTENER_KEY]
+    multicast.register_gateway(entry.data[CONF_HOST], xiaomi_gateway.multicast_callback)
     _LOGGER.debug(
         "Gateway with host '%s' connected, listening for broadcasts",
         entry.data[CONF_HOST],
@@ -201,23 +209,32 @@ 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, config_entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    if entry.data[CONF_KEY] is not None:
+    if config_entry.data[CONF_KEY] is not None:
         platforms = GATEWAY_PLATFORMS
     else:
         platforms = GATEWAY_PLATFORMS_NO_KEY
 
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, platforms)
+    unload_ok = await hass.config_entries.async_unload_platforms(
+        config_entry, platforms
+    )
     if unload_ok:
-        hass.data[DOMAIN][GATEWAYS_KEY].pop(entry.entry_id)
+        hass.data[DOMAIN][GATEWAYS_KEY].pop(config_entry.entry_id)
 
-    if len(hass.data[DOMAIN][GATEWAYS_KEY]) == 0:
+    loaded_entries = [
+        entry
+        for entry in hass.config_entries.async_entries(DOMAIN)
+        if entry.state == ConfigEntryState.LOADED
+    ]
+    if len(loaded_entries) == 1:
         # No gateways left, stop Xiaomi socket
+        unsub_stop = hass.data[DOMAIN].pop(KEY_UNSUB_STOP)
+        unsub_stop()
         hass.data[DOMAIN].pop(GATEWAYS_KEY)
         _LOGGER.debug("Shutting down Xiaomi Gateway Listener")
-        gateway_discovery = hass.data[DOMAIN].pop(LISTENER_KEY)
+        multicast = hass.data[DOMAIN].pop(LISTENER_KEY)
-        await hass.async_add_executor_job(gateway_discovery.stop_listen)
+        multicast.stop_listen()
 
     return unload_ok
 
@@ -262,12 +279,9 @@ class XiaomiDevice(Entity):
             self._is_gateway = False
             self._device_id = self._sid
 
-    def _add_push_data_job(self, *args):
-        self.hass.add_job(self.push_data, *args)
-
     async def async_added_to_hass(self):
         """Start unavailability tracking."""
-        self._xiaomi_hub.callbacks[self._sid].append(self._add_push_data_job)
+        self._xiaomi_hub.callbacks[self._sid].append(self.push_data)
         self._async_track_unavailable()
 
     @property

homeassistant/components/xiaomi_aqara/binary_sensor.py

@@ -287,7 +287,7 @@ class XiaomiMotionSensor(XiaomiBinarySensor):
                 )
 
             if self.entity_id is not None:
-                self._hass.bus.fire(
+                self._hass.bus.async_fire(
                     "xiaomi_aqara.motion", {"entity_id": self.entity_id}
                 )
 
@@ -473,7 +473,7 @@ class XiaomiVibration(XiaomiBinarySensor):
             _LOGGER.warning("Unsupported movement_type detected: %s", value)
             return False
 
-        self.hass.bus.fire(
+        self.hass.bus.async_fire(
             "xiaomi_aqara.movement",
             {"entity_id": self.entity_id, "movement_type": value},
         )
@@ -533,7 +533,7 @@ class XiaomiButton(XiaomiBinarySensor):
             _LOGGER.warning("Unsupported click_type detected: %s", value)
             return False
 
-        self._hass.bus.fire(
+        self._hass.bus.async_fire(
             "xiaomi_aqara.click",
             {"entity_id": self.entity_id, "click_type": click_type},
         )
@@ -570,7 +570,7 @@ class XiaomiCube(XiaomiBinarySensor):
     def parse_data(self, data, raw_data):
         """Parse data sent by gateway."""
         if self._data_key in data:
-            self._hass.bus.fire(
+            self._hass.bus.async_fire(
                 "xiaomi_aqara.cube_action",
                 {"entity_id": self.entity_id, "action_type": data[self._data_key]},
             )
@@ -582,7 +582,7 @@ class XiaomiCube(XiaomiBinarySensor):
                 if isinstance(data["rotate"], int)
                 else data["rotate"].replace(",", ".")
             )
-            self._hass.bus.fire(
+            self._hass.bus.async_fire(
                 "xiaomi_aqara.cube_action",
                 {
                     "entity_id": self.entity_id,
@@ -598,7 +598,7 @@ class XiaomiCube(XiaomiBinarySensor):
                 if isinstance(data["rotate_degree"], int)
                 else data["rotate_degree"].replace(",", ".")
             )
-            self._hass.bus.fire(
+            self._hass.bus.async_fire(
                 "xiaomi_aqara.cube_action",
                 {
                     "entity_id": self.entity_id,

homeassistant/components/xiaomi_aqara/config_flow.py

@@ -106,7 +106,7 @@ class XiaomiAqaraFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
             return await self.async_step_settings()
 
         # Discover Xiaomi Aqara Gateways in the netwerk to get required SIDs.
-        xiaomi = XiaomiGatewayDiscovery(self.hass.add_job, [], self.interface)
+        xiaomi = XiaomiGatewayDiscovery(self.interface)
         try:
             await self.hass.async_add_executor_job(xiaomi.discover_gateways)
         except gaierror:

homeassistant/components/xiaomi_aqara/const.py

@@ -4,6 +4,8 @@ DOMAIN = "xiaomi_aqara"
 
 GATEWAYS_KEY = "gateways"
 LISTENER_KEY = "listener"
+KEY_UNSUB_STOP = "unsub_stop"
+KEY_SETUP_LOCK = "setup_lock"
 
 ZEROCONF_GATEWAY = "lumi-gateway"
 ZEROCONF_ACPARTNER = "lumi-acpartner"

homeassistant/components/xiaomi_aqara/manifest.json

@@ -3,7 +3,7 @@
   "name": "Xiaomi Gateway (Aqara)",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/xiaomi_aqara",
-  "requirements": ["PyXiaomiGateway==0.13.4"],
+  "requirements": ["PyXiaomiGateway==0.14.1"],
   "after_dependencies": ["discovery"],
   "codeowners": ["@danielhiversen", "@syssi"],
   "zeroconf": ["_miio._udp.local."],

requirements_all.txt

@@ -50,7 +50,7 @@ PyTurboJPEG==1.6.7
 PyViCare==2.17.0
 
 # homeassistant.components.xiaomi_aqara
-PyXiaomiGateway==0.13.4
+PyXiaomiGateway==0.14.1
 
 # homeassistant.components.remember_the_milk
 RtmAPI==0.7.2

requirements_test_all.txt

@@ -46,7 +46,7 @@ PyTurboJPEG==1.6.7
 PyViCare==2.17.0
 
 # homeassistant.components.xiaomi_aqara
-PyXiaomiGateway==0.13.4
+PyXiaomiGateway==0.14.1
 
 # homeassistant.components.remember_the_milk
 RtmAPI==0.7.2