Use asyncio in XiaomiAqara instead of threading (#74979)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
2025-04-23 16:57:53 +00:00 · 2022-09-15 08:36:56 +02:00 · 2022-09-15 08:36:56 +02:00 · ade4fcaebd
commit ade4fcaebd
parent 84a812ad05
7 changed files with 53 additions and 37 deletions
--- a/homeassistant/components/xiaomi_aqara/init.py
+++ b/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(
-        LISTENER_KEY,
-        XiaomiGatewayDiscovery(hass.add_job, [], entry.data[CONF_INTERFACE]),
-    )
+    async with setup_lock:
+        if LISTENER_KEY not in hass.data[DOMAIN]:
+            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)
+            # start listining for local pushes (only once)
+            await multicast.start_listen()

-        # register stop callback to shutdown listining for local pushes
-        def stop_xiaomi(event):
-            """Stop Xiaomi Socket."""
-            _LOGGER.debug("Shutting down Xiaomi Gateway Listener")
-            gateway_discovery.stop_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")
+                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)
-        await hass.async_add_executor_job(gateway_discovery.stop_listen)
+        multicast = hass.data[DOMAIN].pop(LISTENER_KEY)
+        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
--- a/homeassistant/components/xiaomi_aqara/binary_sensor.py
+++ b/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,
--- a/homeassistant/components/xiaomi_aqara/config_flow.py
+++ b/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:
--- a/homeassistant/components/xiaomi_aqara/const.py
+++ b/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"
--- a/homeassistant/components/xiaomi_aqara/manifest.json
+++ b/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."],
--- a/requirements_all.txt
+++ b/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
--- a/requirements_test_all.txt
+++ b/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