Move xiaomi_aqara base entity to separate module (#126197)

homeassistant/components/xiaomi_aqara/__init__.py

@@ -1,9 +1,7 @@
 """Support for Xiaomi Gateways."""
 
 import asyncio
-from datetime import timedelta
 import logging
-from typing import Any
 
 import voluptuous as vol
 from xiaomi_gateway import AsyncXiaomiGatewayMulticast, XiaomiGateway
@@ -11,11 +9,8 @@ from xiaomi_gateway import AsyncXiaomiGatewayMulticast, XiaomiGateway
 from homeassistant.components import persistent_notification
 from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import (
-    ATTR_BATTERY_LEVEL,
     ATTR_DEVICE_ID,
-    ATTR_VOLTAGE,
     CONF_HOST,
-    CONF_MAC,
     CONF_PORT,
     CONF_PROTOCOL,
     EVENT_HOMEASSISTANT_STOP,
@@ -24,11 +19,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, ServiceCall, callback
 from homeassistant.helpers import device_registry as dr
 import homeassistant.helpers.config_validation as cv
-from homeassistant.helpers.device_registry import DeviceInfo, format_mac
-from homeassistant.helpers.entity import Entity
-from homeassistant.helpers.event import async_track_point_in_utc_time
 from homeassistant.helpers.typing import ConfigType
-from homeassistant.util.dt import utcnow
 
 from .const import (
     CONF_INTERFACE,
@@ -58,8 +49,6 @@ ATTR_GW_MAC = "gw_mac"
 ATTR_RINGTONE_ID = "ringtone_id"
 ATTR_RINGTONE_VOL = "ringtone_vol"
 
-TIME_TILL_UNAVAILABLE = timedelta(minutes=150)
-
 SERVICE_PLAY_RINGTONE = "play_ringtone"
 SERVICE_STOP_RINGTONE = "stop_ringtone"
 SERVICE_ADD_DEVICE = "add_device"
@@ -245,152 +234,6 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
     return unload_ok
 
 
-class XiaomiDevice(Entity):
-    """Representation a base Xiaomi device."""
-
-    _attr_should_poll = False
-
-    def __init__(self, device, device_type, xiaomi_hub, config_entry):
-        """Initialize the Xiaomi device."""
-        self._state = None
-        self._is_available = True
-        self._sid = device["sid"]
-        self._model = device["model"]
-        self._protocol = device["proto"]
-        self._name = f"{device_type}_{self._sid}"
-        self._device_name = f"{self._model}_{self._sid}"
-        self._type = device_type
-        self._write_to_hub = xiaomi_hub.write_to_hub
-        self._get_from_hub = xiaomi_hub.get_from_hub
-        self._extra_state_attributes = {}
-        self._remove_unavailability_tracker = None
-        self._xiaomi_hub = xiaomi_hub
-        self.parse_data(device["data"], device["raw_data"])
-        self.parse_voltage(device["data"])
-
-        if hasattr(self, "_data_key") and self._data_key:
-            self._unique_id = f"{self._data_key}{self._sid}"
-        else:
-            self._unique_id = f"{self._type}{self._sid}"
-
-        self._gateway_id = config_entry.unique_id
-        if config_entry.data[CONF_MAC] == format_mac(self._sid):
-            # this entity belongs to the gateway itself
-            self._is_gateway = True
-            self._device_id = config_entry.unique_id
-        else:
-            # this entity is connected through zigbee
-            self._is_gateway = False
-            self._device_id = self._sid
-
-    async def async_added_to_hass(self):
-        """Start unavailability tracking."""
-        self._xiaomi_hub.callbacks[self._sid].append(self.push_data)
-        self._async_track_unavailable()
-
-    @property
-    def name(self):
-        """Return the name of the device."""
-        return self._name
-
-    @property
-    def unique_id(self) -> str:
-        """Return a unique ID."""
-        return self._unique_id
-
-    @property
-    def device_id(self):
-        """Return the device id of the Xiaomi Aqara device."""
-        return self._device_id
-
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return the device info of the Xiaomi Aqara device."""
-        if self._is_gateway:
-            device_info = DeviceInfo(
-                identifiers={(DOMAIN, self._device_id)},
-                model=self._model,
-            )
-        else:
-            device_info = DeviceInfo(
-                connections={(dr.CONNECTION_ZIGBEE, self._device_id)},
-                identifiers={(DOMAIN, self._device_id)},
-                manufacturer="Xiaomi Aqara",
-                model=self._model,
-                name=self._device_name,
-                sw_version=self._protocol,
-                via_device=(DOMAIN, self._gateway_id),
-            )
-
-        return device_info
-
-    @property
-    def available(self):
-        """Return True if entity is available."""
-        return self._is_available
-
-    @property
-    def extra_state_attributes(self):
-        """Return the state attributes."""
-        return self._extra_state_attributes
-
-    @callback
-    def _async_set_unavailable(self, now):
-        """Set state to UNAVAILABLE."""
-        self._remove_unavailability_tracker = None
-        self._is_available = False
-        self.async_write_ha_state()
-
-    @callback
-    def _async_track_unavailable(self):
-        if self._remove_unavailability_tracker:
-            self._remove_unavailability_tracker()
-        self._remove_unavailability_tracker = async_track_point_in_utc_time(
-            self.hass, self._async_set_unavailable, utcnow() + TIME_TILL_UNAVAILABLE
-        )
-        if not self._is_available:
-            self._is_available = True
-            return True
-        return False
-
-    def push_data(self, data: dict[str, Any], raw_data: dict[Any, Any]) -> None:
-        """Push from Hub running in another thread."""
-        self.hass.loop.call_soon_threadsafe(self.async_push_data, data, raw_data)
-
-    @callback
-    def async_push_data(self, data: dict[str, Any], raw_data: dict[Any, Any]) -> None:
-        """Push from Hub handled in the event loop."""
-        _LOGGER.debug("PUSH >> %s: %s", self, data)
-        was_unavailable = self._async_track_unavailable()
-        is_data = self.parse_data(data, raw_data)
-        is_voltage = self.parse_voltage(data)
-        if is_data or is_voltage or was_unavailable:
-            self.async_write_ha_state()
-
-    def parse_voltage(self, data):
-        """Parse battery level data sent by gateway."""
-        if "voltage" in data:
-            voltage_key = "voltage"
-        elif "battery_voltage" in data:
-            voltage_key = "battery_voltage"
-        else:
-            return False
-
-        max_volt = 3300
-        min_volt = 2800
-        voltage = data[voltage_key]
-        self._extra_state_attributes[ATTR_VOLTAGE] = round(voltage / 1000.0, 2)
-        voltage = min(voltage, max_volt)
-        voltage = max(voltage, min_volt)
-        percent = ((voltage - min_volt) / (max_volt - min_volt)) * 100
-        self._extra_state_attributes[ATTR_BATTERY_LEVEL] = round(percent, 1)
-        return True
-
-    def parse_data(self, data, raw_data):
-        """Parse data sent by gateway."""
-        raise NotImplementedError
-
-
 def _add_gateway_to_schema(hass, schema):
     """Extend a voluptuous schema with a gateway validator."""
 

homeassistant/components/xiaomi_aqara/binary_sensor.py

@@ -12,8 +12,8 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_call_later
 from homeassistant.helpers.restore_state import RestoreEntity
 
-from . import XiaomiDevice
 from .const import DOMAIN, GATEWAYS_KEY
+from .entity import XiaomiDevice
 
 _LOGGER = logging.getLogger(__name__)
 

homeassistant/components/xiaomi_aqara/cover.py

@@ -7,8 +7,8 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import XiaomiDevice
 from .const import DOMAIN, GATEWAYS_KEY
+from .entity import XiaomiDevice
 
 ATTR_CURTAIN_LEVEL = "curtain_level"
 

homeassistant/components/xiaomi_aqara/entity.py

@@ -0,0 +1,165 @@
+"""Support for Xiaomi Gateways."""
+
+from datetime import timedelta
+import logging
+from typing import Any
+
+from homeassistant.const import ATTR_BATTERY_LEVEL, ATTR_VOLTAGE, CONF_MAC
+from homeassistant.core import callback
+from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers.device_registry import DeviceInfo, format_mac
+from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.event import async_track_point_in_utc_time
+from homeassistant.util.dt import utcnow
+
+from .const import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+TIME_TILL_UNAVAILABLE = timedelta(minutes=150)
+
+
+class XiaomiDevice(Entity):
+    """Representation a base Xiaomi device."""
+
+    _attr_should_poll = False
+
+    def __init__(self, device, device_type, xiaomi_hub, config_entry):
+        """Initialize the Xiaomi device."""
+        self._state = None
+        self._is_available = True
+        self._sid = device["sid"]
+        self._model = device["model"]
+        self._protocol = device["proto"]
+        self._name = f"{device_type}_{self._sid}"
+        self._device_name = f"{self._model}_{self._sid}"
+        self._type = device_type
+        self._write_to_hub = xiaomi_hub.write_to_hub
+        self._get_from_hub = xiaomi_hub.get_from_hub
+        self._extra_state_attributes = {}
+        self._remove_unavailability_tracker = None
+        self._xiaomi_hub = xiaomi_hub
+        self.parse_data(device["data"], device["raw_data"])
+        self.parse_voltage(device["data"])
+
+        if hasattr(self, "_data_key") and self._data_key:
+            self._unique_id = f"{self._data_key}{self._sid}"
+        else:
+            self._unique_id = f"{self._type}{self._sid}"
+
+        self._gateway_id = config_entry.unique_id
+        if config_entry.data[CONF_MAC] == format_mac(self._sid):
+            # this entity belongs to the gateway itself
+            self._is_gateway = True
+            self._device_id = config_entry.unique_id
+        else:
+            # this entity is connected through zigbee
+            self._is_gateway = False
+            self._device_id = self._sid
+
+    async def async_added_to_hass(self):
+        """Start unavailability tracking."""
+        self._xiaomi_hub.callbacks[self._sid].append(self.push_data)
+        self._async_track_unavailable()
+
+    @property
+    def name(self):
+        """Return the name of the device."""
+        return self._name
+
+    @property
+    def unique_id(self) -> str:
+        """Return a unique ID."""
+        return self._unique_id
+
+    @property
+    def device_id(self):
+        """Return the device id of the Xiaomi Aqara device."""
+        return self._device_id
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return the device info of the Xiaomi Aqara device."""
+        if self._is_gateway:
+            device_info = DeviceInfo(
+                identifiers={(DOMAIN, self._device_id)},
+                model=self._model,
+            )
+        else:
+            device_info = DeviceInfo(
+                connections={(dr.CONNECTION_ZIGBEE, self._device_id)},
+                identifiers={(DOMAIN, self._device_id)},
+                manufacturer="Xiaomi Aqara",
+                model=self._model,
+                name=self._device_name,
+                sw_version=self._protocol,
+                via_device=(DOMAIN, self._gateway_id),
+            )
+
+        return device_info
+
+    @property
+    def available(self):
+        """Return True if entity is available."""
+        return self._is_available
+
+    @property
+    def extra_state_attributes(self):
+        """Return the state attributes."""
+        return self._extra_state_attributes
+
+    @callback
+    def _async_set_unavailable(self, now):
+        """Set state to UNAVAILABLE."""
+        self._remove_unavailability_tracker = None
+        self._is_available = False
+        self.async_write_ha_state()
+
+    @callback
+    def _async_track_unavailable(self):
+        if self._remove_unavailability_tracker:
+            self._remove_unavailability_tracker()
+        self._remove_unavailability_tracker = async_track_point_in_utc_time(
+            self.hass, self._async_set_unavailable, utcnow() + TIME_TILL_UNAVAILABLE
+        )
+        if not self._is_available:
+            self._is_available = True
+            return True
+        return False
+
+    def push_data(self, data: dict[str, Any], raw_data: dict[Any, Any]) -> None:
+        """Push from Hub running in another thread."""
+        self.hass.loop.call_soon_threadsafe(self.async_push_data, data, raw_data)
+
+    @callback
+    def async_push_data(self, data: dict[str, Any], raw_data: dict[Any, Any]) -> None:
+        """Push from Hub handled in the event loop."""
+        _LOGGER.debug("PUSH >> %s: %s", self, data)
+        was_unavailable = self._async_track_unavailable()
+        is_data = self.parse_data(data, raw_data)
+        is_voltage = self.parse_voltage(data)
+        if is_data or is_voltage or was_unavailable:
+            self.async_write_ha_state()
+
+    def parse_voltage(self, data):
+        """Parse battery level data sent by gateway."""
+        if "voltage" in data:
+            voltage_key = "voltage"
+        elif "battery_voltage" in data:
+            voltage_key = "battery_voltage"
+        else:
+            return False
+
+        max_volt = 3300
+        min_volt = 2800
+        voltage = data[voltage_key]
+        self._extra_state_attributes[ATTR_VOLTAGE] = round(voltage / 1000.0, 2)
+        voltage = min(voltage, max_volt)
+        voltage = max(voltage, min_volt)
+        percent = ((voltage - min_volt) / (max_volt - min_volt)) * 100
+        self._extra_state_attributes[ATTR_BATTERY_LEVEL] = round(percent, 1)
+        return True
+
+    def parse_data(self, data, raw_data):
+        """Parse data sent by gateway."""
+        raise NotImplementedError

homeassistant/components/xiaomi_aqara/light.py

@@ -16,8 +16,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.color as color_util
 
-from . import XiaomiDevice
 from .const import DOMAIN, GATEWAYS_KEY
+from .entity import XiaomiDevice
 
 _LOGGER = logging.getLogger(__name__)
 

homeassistant/components/xiaomi_aqara/lock.py

@@ -9,8 +9,8 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_call_later
 
-from . import XiaomiDevice
 from .const import DOMAIN, GATEWAYS_KEY
+from .entity import XiaomiDevice
 
 FINGER_KEY = "fing_verified"
 PASSWORD_KEY = "psw_verified"

homeassistant/components/xiaomi_aqara/sensor.py

@@ -22,8 +22,8 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import XiaomiDevice
 from .const import BATTERY_MODELS, DOMAIN, GATEWAYS_KEY, POWER_MODELS
+from .entity import XiaomiDevice
 
 _LOGGER = logging.getLogger(__name__)
 

homeassistant/components/xiaomi_aqara/switch.py

@@ -8,8 +8,8 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import XiaomiDevice
 from .const import DOMAIN, GATEWAYS_KEY
+from .entity import XiaomiDevice
 
 _LOGGER = logging.getLogger(__name__)