Use registry to find linked batteries for homekit

homeassistant/components/homekit/__init__.py

@@ -7,10 +7,13 @@
 from zeroconf import InterfaceChoice
 
 from homeassistant.components import cover, vacuum
+from homeassistant.components.binary_sensor import DEVICE_CLASS_BATTERY_CHARGING
 from homeassistant.components.cover import DEVICE_CLASS_GARAGE, DEVICE_CLASS_GATE
 from homeassistant.components.http import HomeAssistantView
 from homeassistant.components.media_player import DEVICE_CLASS_TV
 from homeassistant.const import (
+    ATTR_BATTERY_CHARGING,
+    ATTR_BATTERY_LEVEL,
     ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
     ATTR_SERVICE,
@@ -20,6 +23,7 @@
     CONF_NAME,
     CONF_PORT,
     CONF_TYPE,
+    DEVICE_CLASS_BATTERY,
     DEVICE_CLASS_HUMIDITY,
     DEVICE_CLASS_ILLUMINANCE,
     DEVICE_CLASS_TEMPERATURE,
@@ -31,6 +35,7 @@
 )
 from homeassistant.core import callback
 from homeassistant.exceptions import Unauthorized
+from homeassistant.helpers import entity_registry
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entityfilter import FILTER_SCHEMA
 from homeassistant.util import get_local_ip
@@ -47,6 +52,8 @@
     CONF_ENTITY_CONFIG,
     CONF_FEATURE_LIST,
     CONF_FILTER,
+    CONF_LINKED_BATTERY_CHARGING_SENSOR,
+    CONF_LINKED_BATTERY_SENSOR,
     CONF_SAFE_MODE,
     CONF_ZEROCONF_DEFAULT_INTERFACE,
     DEFAULT_AUTO_START,
@@ -202,21 +209,21 @@ def async_describe_logbook_event(event):
     )
 
     if auto_start:
-        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_START, homekit.start)
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_START, homekit.async_start)
         return True
 
-    def handle_homekit_service_start(service):
+    async def async_handle_homekit_service_start(service):
         """Handle start HomeKit service call."""
         if homekit.status != STATUS_READY:
             _LOGGER.warning(
                 "HomeKit is not ready. Either it is already running or has "
                 "been stopped."
             )
             return
-        homekit.start()
+        await homekit.async_start()
 
     hass.services.async_register(
-        DOMAIN, SERVICE_HOMEKIT_START, handle_homekit_service_start
+        DOMAIN, SERVICE_HOMEKIT_START, async_handle_homekit_service_start
     )
 
     return True
@@ -355,7 +362,7 @@ def setup(self):
         # pylint: disable=import-outside-toplevel
         from .accessories import HomeBridge, HomeDriver
 
-        self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, self.stop)
+        self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, self.async_stop)
 
         ip_addr = self._ip_address or get_local_ip()
         path = self.hass.config.path(HOMEKIT_FILE)
@@ -393,7 +400,7 @@ def reset_accessories(self, entity_ids):
 
     def add_bridge_accessory(self, state):
         """Try adding accessory to bridge if configured beforehand."""
-        if not state or not self._filter(state.entity_id):
+        if not self._filter(state.entity_id):
             return
 
         # The bridge itself counts as an accessory
@@ -428,12 +435,32 @@ def remove_bridge_accessory(self, aid):
             acc = self.bridge.accessories.pop(aid)
         return acc
 
-    def start(self, *args):
+    async def async_start(self, *args):
         """Start the accessory driver."""
         if self.status != STATUS_READY:
             return
         self.status = STATUS_WAIT
 
+        ent_reg = await entity_registry.async_get_registry(self.hass)
+
+        device_lookup = ent_reg.async_get_device_class_lookup(
+            {
+                ("binary_sensor", DEVICE_CLASS_BATTERY_CHARGING),
+                ("sensor", DEVICE_CLASS_BATTERY),
+            }
+        )
+
+        bridged_states = []
+        for state in self.hass.states.async_all():
+            if not self._filter(state.entity_id):
+                continue
+
+            self._async_configure_linked_battery_sensors(ent_reg, device_lookup, state)
+            bridged_states.append(state)
+
+        await self.hass.async_add_executor_job(self._start, bridged_states)
+
+    def _start(self, bridged_states):
         from . import (  # noqa: F401 pylint: disable=unused-import, import-outside-toplevel
             type_covers,
             type_fans,
@@ -446,7 +473,7 @@ def start(self, *args):
             type_thermostats,
         )
 
-        for state in self.hass.states.all():
+        for state in bridged_states:
             self.add_bridge_accessory(state)
 
         self.driver.add_accessory(self.bridge)
@@ -457,17 +484,49 @@ def start(self, *args):
             )
 
         _LOGGER.debug("Driver start")
-        self.hass.add_job(self.driver.start)
+        self.hass.async_add_executor_job(self.driver.start)
         self.status = STATUS_RUNNING
 
-    def stop(self, *args):
+    async def async_stop(self, *args):
         """Stop the accessory driver."""
         if self.status != STATUS_RUNNING:
             return
         self.status = STATUS_STOPPED
 
         _LOGGER.debug("Driver stop")
-        self.hass.add_job(self.driver.stop)
+        self.hass.async_add_executor_job(self.driver.stop)
+
+    @callback
+    def _async_configure_linked_battery_sensors(self, ent_reg, device_lookup, state):
+        entry = ent_reg.async_get(state.entity_id)
+
+        if (
+            entry is None
+            or entry.device_id is None
+            or entry.device_id not in device_lookup
+            or entry.device_class
+            in (DEVICE_CLASS_BATTERY_CHARGING, DEVICE_CLASS_BATTERY)
+        ):
+            return
+
+        if ATTR_BATTERY_CHARGING not in state.attributes:
+            battery_charging_binary_sensor_entity_id = device_lookup[
+                entry.device_id
+            ].get(("binary_sensor", DEVICE_CLASS_BATTERY_CHARGING))
+            if battery_charging_binary_sensor_entity_id:
+                self._config.setdefault(state.entity_id, {}).setdefault(
+                    CONF_LINKED_BATTERY_CHARGING_SENSOR,
+                    battery_charging_binary_sensor_entity_id,
+                )
+
+        if ATTR_BATTERY_LEVEL not in state.attributes:
+            battery_sensor_entity_id = device_lookup[entry.device_id].get(
+                ("sensor", DEVICE_CLASS_BATTERY)
+            )
+            if battery_sensor_entity_id:
+                self._config.setdefault(state.entity_id, {}).setdefault(
+                    CONF_LINKED_BATTERY_SENSOR, battery_sensor_entity_id
+                )
 
 
 class HomeKitPairingQRView(HomeAssistantView):

homeassistant/components/homekit/accessories.py

@@ -13,6 +13,7 @@
     ATTR_BATTERY_LEVEL,
     ATTR_ENTITY_ID,
     ATTR_SERVICE,
+    STATE_ON,
     __version__,
 )
 from homeassistant.core import callback as ha_callback, split_entity_id
@@ -30,11 +31,15 @@
     CHAR_BATTERY_LEVEL,
     CHAR_CHARGING_STATE,
     CHAR_STATUS_LOW_BATTERY,
+    CONF_LINKED_BATTERY_CHARGING_SENSOR,
     CONF_LINKED_BATTERY_SENSOR,
     CONF_LOW_BATTERY_THRESHOLD,
     DEBOUNCE_TIMEOUT,
     DEFAULT_LOW_BATTERY_THRESHOLD,
     EVENT_HOMEKIT_CHANGED,
+    HK_CHARGING,
+    HK_NOT_CHARGABLE,
+    HK_NOT_CHARGING,
     MANUFACTURER,
     SERV_BATTERY_SERVICE,
 )
@@ -94,17 +99,17 @@ def __init__(
         self.entity_id = entity_id
         self.hass = hass
         self.debounce = {}
-        self._support_battery_level = False
-        self._support_battery_charging = True
         self.linked_battery_sensor = self.config.get(CONF_LINKED_BATTERY_SENSOR)
+        self.linked_battery_charging_sensor = self.config.get(
+            CONF_LINKED_BATTERY_CHARGING_SENSOR
+        )
         self.low_battery_threshold = self.config.get(
             CONF_LOW_BATTERY_THRESHOLD, DEFAULT_LOW_BATTERY_THRESHOLD
         )
 
         """Add battery service if available"""
-        battery_found = self.hass.states.get(self.entity_id).attributes.get(
-            ATTR_BATTERY_LEVEL
-        )
+        entity_attributes = self.hass.states.get(self.entity_id).attributes
+        battery_found = entity_attributes.get(ATTR_BATTERY_LEVEL)
 
         if self.linked_battery_sensor:
             state = self.hass.states.get(self.linked_battery_sensor)
@@ -118,13 +123,28 @@ def __init__(
                     self.linked_battery_sensor,
                 )
 
-        if battery_found is None:
+        if not battery_found:
             return
+
         _LOGGER.debug("%s: Found battery level", self.entity_id)
-        self._support_battery_level = True
+
+        if self.linked_battery_charging_sensor:
+            state = self.hass.states.get(self.linked_battery_charging_sensor)
+            if state is None:
+                self.linked_battery_charging_sensor = None
+                _LOGGER.warning(
+                    "%s: Battery charging binary_sensor state missing: %s",
+                    self.entity_id,
+                    self.linked_battery_charging_sensor,
+                )
+            else:
+                _LOGGER.debug("%s: Found battery charging", self.entity_id)
+
         serv_battery = self.add_preload_service(SERV_BATTERY_SERVICE)
         self._char_battery = serv_battery.configure_char(CHAR_BATTERY_LEVEL, value=0)
-        self._char_charging = serv_battery.configure_char(CHAR_CHARGING_STATE, value=2)
+        self._char_charging = serv_battery.configure_char(
+            CHAR_CHARGING_STATE, value=HK_NOT_CHARGABLE
+        )
         self._char_low_battery = serv_battery.configure_char(
             CHAR_STATUS_LOW_BATTERY, value=0
         )
@@ -142,55 +162,111 @@ async def run_handler(self):
         Run inside the Home Assistant event loop.
         """
         state = self.hass.states.get(self.entity_id)
-        self.hass.async_add_job(self.update_state_callback, None, None, state)
+        self.hass.async_add_executor_job(self.update_state_callback, None, None, state)
         async_track_state_change(self.hass, self.entity_id, self.update_state_callback)
 
+        battery_charging_state = None
+        battery_state = None
         if self.linked_battery_sensor:
-            battery_state = self.hass.states.get(self.linked_battery_sensor)
-            self.hass.async_add_job(
-                self.update_linked_battery, None, None, battery_state
+            linked_battery_sensor_state = self.hass.states.get(
+                self.linked_battery_sensor
+            )
+            battery_state = linked_battery_sensor_state.state
+            battery_charging_state = linked_battery_sensor_state.attributes.get(
+                ATTR_BATTERY_CHARGING
             )
             async_track_state_change(
                 self.hass, self.linked_battery_sensor, self.update_linked_battery
             )
+        else:
+            battery_state = state.attributes.get(ATTR_BATTERY_LEVEL)
+        if self.linked_battery_charging_sensor:
+            battery_charging_state = (
+                self.hass.states.get(self.linked_battery_charging_sensor).state
+                == STATE_ON
+            )
+            async_track_state_change(
+                self.hass,
+                self.linked_battery_charging_sensor,
+                self.update_linked_battery_charging,
+            )
+        elif battery_charging_state is None:
+            battery_charging_state = state.attributes.get(ATTR_BATTERY_CHARGING)
+
+        if battery_state is not None or battery_charging_state is not None:
+            self.hass.async_add_executor_job(
+                self.update_battery, battery_state, battery_charging_state
+            )
 
     @ha_callback
     def update_state_callback(self, entity_id=None, old_state=None, new_state=None):
         """Handle state change listener callback."""
         _LOGGER.debug("New_state: %s", new_state)
         if new_state is None:
             return
-        if self._support_battery_level and not self.linked_battery_sensor:
-            self.hass.async_add_executor_job(self.update_battery, new_state)
+        battery_state = None
+        battery_charging_state = None
+        if (
+            not self.linked_battery_sensor
+            and ATTR_BATTERY_LEVEL in new_state.attributes
+        ):
+            battery_state = new_state.attributes.get(ATTR_BATTERY_LEVEL)
+        if (
+            not self.linked_battery_charging_sensor
+            and ATTR_BATTERY_CHARGING in new_state.attributes
+        ):
+            battery_charging_state = new_state.attributes.get(ATTR_BATTERY_CHARGING)
+        if battery_state is not None or battery_charging_state is not None:
+            self.hass.async_add_executor_job(
+                self.update_battery, battery_state, battery_charging_state
+            )
         self.hass.async_add_executor_job(self.update_state, new_state)
 
     @ha_callback
     def update_linked_battery(self, entity_id=None, old_state=None, new_state=None):
         """Handle linked battery sensor state change listener callback."""
-        self.hass.async_add_executor_job(self.update_battery, new_state)
+        if self.linked_battery_charging_sensor:
+            battery_charging_state = None
+        else:
+            battery_charging_state = new_state.attributes.get(ATTR_BATTERY_CHARGING)
+        self.hass.async_add_executor_job(
+            self.update_battery, new_state.state, battery_charging_state,
+        )
+
+    @ha_callback
+    def update_linked_battery_charging(
+        self, entity_id=None, old_state=None, new_state=None
+    ):
+        """Handle linked battery charging sensor state change listener callback."""
+        self.hass.async_add_executor_job(
+            self.update_battery, None, new_state.state == STATE_ON
+        )
 
-    def update_battery(self, new_state):
+    def update_battery(self, battery_level, battery_charging):
         """Update battery service if available.
 
         Only call this function if self._support_battery_level is True.
         """
-        battery_level = convert_to_float(new_state.attributes.get(ATTR_BATTERY_LEVEL))
-        if self.linked_battery_sensor:
-            battery_level = convert_to_float(new_state.state)
-        if battery_level is None:
-            return
-        self._char_battery.set_value(battery_level)
-        self._char_low_battery.set_value(battery_level < self.low_battery_threshold)
-        _LOGGER.debug("%s: Updated battery level to %d", self.entity_id, battery_level)
-        if not self._support_battery_charging:
-            return
-        charging = new_state.attributes.get(ATTR_BATTERY_CHARGING)
-        if charging is None:
-            self._support_battery_charging = False
+        battery_level = convert_to_float(battery_level)
+        if battery_level is not None:
+            if self._char_battery.value != battery_level:
+                self._char_battery.set_value(battery_level)
+            is_low_battery = 1 if battery_level < self.low_battery_threshold else 0
+            if self._char_low_battery.value != is_low_battery:
+                self._char_low_battery.set_value(is_low_battery)
+                _LOGGER.debug(
+                    "%s: Updated battery level to %d", self.entity_id, battery_level
+                )
+
+        if battery_charging is None:
             return
-        hk_charging = 1 if charging is True else 0
-        self._char_charging.set_value(hk_charging)
-        _LOGGER.debug("%s: Updated battery charging to %d", self.entity_id, hk_charging)
+
+        hk_charging = HK_CHARGING if battery_charging else HK_NOT_CHARGING
+        if self._char_charging.value != hk_charging:
+            self._char_charging.set_value(hk_charging)
+            _LOGGER.debug(
+                "%s: Updated battery charging to %d", self.entity_id, hk_charging
+            )
 
     def update_state(self, new_state):
         """Handle state change to update HomeKit value.