From 28197adebd13de0a64dbb7ed44ab4b9a44511207 Mon Sep 17 00:00:00 2001
From: Ernst Klamer <e.klamer@gmail.com>
Date: Mon, 24 Jul 2023 21:13:16 +0200
Subject: [PATCH] Add support for sleepy Xiaomi BLE sensors (#97166)

---
 .../components/xiaomi_ble/__init__.py         |  34 ++++-
 .../components/xiaomi_ble/binary_sensor.py    |  25 ++--
 homeassistant/components/xiaomi_ble/const.py  |   4 +-
 .../components/xiaomi_ble/coordinator.py      |  63 +++++++++
 homeassistant/components/xiaomi_ble/sensor.py |  22 +++-
 tests/components/xiaomi_ble/test_sensor.py    | 121 +++++++++++++++++-
 6 files changed, 242 insertions(+), 27 deletions(-)
 create mode 100644 homeassistant/components/xiaomi_ble/coordinator.py

diff --git a/homeassistant/components/xiaomi_ble/__init__.py b/homeassistant/components/xiaomi_ble/__init__.py
index 3930c50c70c..1810d52323c 100644
--- a/homeassistant/components/xiaomi_ble/__init__.py
+++ b/homeassistant/components/xiaomi_ble/__init__.py
@@ -12,15 +12,18 @@ from homeassistant.components.bluetooth import (
     BluetoothServiceInfoBleak,
     async_ble_device_from_address,
 )
-from homeassistant.components.bluetooth.active_update_processor import (
-    ActiveBluetoothProcessorCoordinator,
-)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import CoreState, HomeAssistant
 from homeassistant.helpers.device_registry import DeviceRegistry, async_get
 
-from .const import DOMAIN, XIAOMI_BLE_EVENT, XiaomiBleEvent
+from .const import (
+    CONF_DISCOVERED_EVENT_CLASSES,
+    DOMAIN,
+    XIAOMI_BLE_EVENT,
+    XiaomiBleEvent,
+)
+from .coordinator import XiaomiActiveBluetoothProcessorCoordinator
 
 PLATFORMS: list[Platform] = [Platform.BINARY_SENSOR, Platform.SENSOR]
 
@@ -36,6 +39,10 @@ def process_service_info(
 ) -> SensorUpdate:
     """Process a BluetoothServiceInfoBleak, running side effects and returning sensor data."""
     update = data.update(service_info)
+    coordinator: XiaomiActiveBluetoothProcessorCoordinator = hass.data[DOMAIN][
+        entry.entry_id
+    ]
+    discovered_device_classes = coordinator.discovered_device_classes
     if update.events:
         address = service_info.device.address
         for device_key, event in update.events.items():
@@ -49,6 +56,16 @@ def process_service_info(
                 sw_version=sensor_device_info.sw_version,
                 hw_version=sensor_device_info.hw_version,
             )
+            event_class = event.device_key.key
+            event_type = event.event_type
+
+            if event_class not in discovered_device_classes:
+                discovered_device_classes.add(event_class)
+                hass.config_entries.async_update_entry(
+                    entry,
+                    data=entry.data
+                    | {CONF_DISCOVERED_EVENT_CLASSES: list(discovered_device_classes)},
+                )
 
             hass.bus.async_fire(
                 XIAOMI_BLE_EVENT,
@@ -56,7 +73,8 @@ def process_service_info(
                     XiaomiBleEvent(
                         device_id=device.id,
                         address=address,
-                        event_type=event.event_type,
+                        event_class=event_class,  # ie 'button'
+                        event_type=event_type,  # ie 'press'
                         event_properties=event.event_properties,
                     )
                 ),
@@ -121,7 +139,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     device_registry = async_get(hass)
     coordinator = hass.data.setdefault(DOMAIN, {})[
         entry.entry_id
-    ] = ActiveBluetoothProcessorCoordinator(
+    ] = XiaomiActiveBluetoothProcessorCoordinator(
         hass,
         _LOGGER,
         address=address,
@@ -130,6 +148,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             hass, entry, data, service_info, device_registry
         ),
         needs_poll_method=_needs_poll,
+        device_data=data,
+        discovered_device_classes=set(
+            entry.data.get(CONF_DISCOVERED_EVENT_CLASSES, [])
+        ),
         poll_method=_async_poll,
         # We will take advertisements from non-connectable devices
         # since we will trade the BLEDevice for a connectable one
diff --git a/homeassistant/components/xiaomi_ble/binary_sensor.py b/homeassistant/components/xiaomi_ble/binary_sensor.py
index 3d7bdfd0b48..f7c4c87014c 100644
--- a/homeassistant/components/xiaomi_ble/binary_sensor.py
+++ b/homeassistant/components/xiaomi_ble/binary_sensor.py
@@ -1,7 +1,6 @@
 """Support for Xiaomi binary sensors."""
 from __future__ import annotations
 
-from xiaomi_ble import SLEEPY_DEVICE_MODELS
 from xiaomi_ble.parser import (
     BinarySensorDeviceClass as XiaomiBinarySensorDeviceClass,
     ExtendedBinarySensorDeviceClass,
@@ -15,17 +14,18 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntityDescription,
 )
 from homeassistant.components.bluetooth.passive_update_processor import (
-    PassiveBluetoothDataProcessor,
     PassiveBluetoothDataUpdate,
-    PassiveBluetoothProcessorCoordinator,
     PassiveBluetoothProcessorEntity,
 )
-from homeassistant.const import ATTR_MODEL
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.sensor import sensor_device_info_to_hass_device_info
 
 from .const import DOMAIN
+from .coordinator import (
+    XiaomiActiveBluetoothProcessorCoordinator,
+    XiaomiPassiveBluetoothDataProcessor,
+)
 from .device import device_key_to_bluetooth_entity_key
 
 BINARY_SENSOR_DESCRIPTIONS = {
@@ -108,10 +108,12 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Xiaomi BLE sensors."""
-    coordinator: PassiveBluetoothProcessorCoordinator = hass.data[DOMAIN][
+    coordinator: XiaomiActiveBluetoothProcessorCoordinator = hass.data[DOMAIN][
         entry.entry_id
     ]
-    processor = PassiveBluetoothDataProcessor(sensor_update_to_bluetooth_data_update)
+    processor = XiaomiPassiveBluetoothDataProcessor(
+        sensor_update_to_bluetooth_data_update
+    )
     entry.async_on_unload(
         processor.async_add_entities_listener(
             XiaomiBluetoothSensorEntity, async_add_entities
@@ -121,7 +123,7 @@ async def async_setup_entry(
 
 
 class XiaomiBluetoothSensorEntity(
-    PassiveBluetoothProcessorEntity[PassiveBluetoothDataProcessor[bool | None]],
+    PassiveBluetoothProcessorEntity[XiaomiPassiveBluetoothDataProcessor],
     BinarySensorEntity,
 ):
     """Representation of a Xiaomi binary sensor."""
@@ -134,8 +136,7 @@ class XiaomiBluetoothSensorEntity(
     @property
     def available(self) -> bool:
         """Return True if entity is available."""
-        if self.device_info and self.device_info[ATTR_MODEL] in SLEEPY_DEVICE_MODELS:
-            # These devices sleep for an indeterminate amount of time
-            # so there is no way to track their availability.
-            return True
-        return super().available
+        coordinator: XiaomiActiveBluetoothProcessorCoordinator = (
+            self.processor.coordinator
+        )
+        return coordinator.device_data.sleepy_device or super().available
diff --git a/homeassistant/components/xiaomi_ble/const.py b/homeassistant/components/xiaomi_ble/const.py
index dda6c61d8aa..1566478bcea 100644
--- a/homeassistant/components/xiaomi_ble/const.py
+++ b/homeassistant/components/xiaomi_ble/const.py
@@ -6,6 +6,7 @@ from typing import Final, TypedDict
 DOMAIN = "xiaomi_ble"
 
 
+CONF_DISCOVERED_EVENT_CLASSES: Final = "known_events"
 CONF_EVENT_PROPERTIES: Final = "event_properties"
 EVENT_PROPERTIES: Final = "event_properties"
 EVENT_TYPE: Final = "event_type"
@@ -17,5 +18,6 @@ class XiaomiBleEvent(TypedDict):
 
     device_id: str
     address: str
-    event_type: str
+    event_class: str  # ie 'button'
+    event_type: str  # ie 'press'
     event_properties: dict[str, str | int | float | None] | None
diff --git a/homeassistant/components/xiaomi_ble/coordinator.py b/homeassistant/components/xiaomi_ble/coordinator.py
new file mode 100644
index 00000000000..2a4b35f6171
--- /dev/null
+++ b/homeassistant/components/xiaomi_ble/coordinator.py
@@ -0,0 +1,63 @@
+"""The Xiaomi BLE integration."""
+from collections.abc import Callable, Coroutine
+from logging import Logger
+from typing import Any
+
+from xiaomi_ble import XiaomiBluetoothDeviceData
+
+from homeassistant.components.bluetooth import (
+    BluetoothScanningMode,
+    BluetoothServiceInfoBleak,
+)
+from homeassistant.components.bluetooth.active_update_processor import (
+    ActiveBluetoothProcessorCoordinator,
+)
+from homeassistant.components.bluetooth.passive_update_processor import (
+    PassiveBluetoothDataProcessor,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.debounce import Debouncer
+
+
+class XiaomiActiveBluetoothProcessorCoordinator(ActiveBluetoothProcessorCoordinator):
+    """Define a Xiaomi Bluetooth Active Update Processor Coordinator."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        logger: Logger,
+        *,
+        address: str,
+        mode: BluetoothScanningMode,
+        update_method: Callable[[BluetoothServiceInfoBleak], Any],
+        needs_poll_method: Callable[[BluetoothServiceInfoBleak, float | None], bool],
+        device_data: XiaomiBluetoothDeviceData,
+        discovered_device_classes: set[str],
+        poll_method: Callable[
+            [BluetoothServiceInfoBleak],
+            Coroutine[Any, Any, Any],
+        ]
+        | None = None,
+        poll_debouncer: Debouncer[Coroutine[Any, Any, None]] | None = None,
+        connectable: bool = True,
+    ) -> None:
+        """Initialize the Xiaomi Bluetooth Active Update Processor Coordinator."""
+        super().__init__(
+            hass=hass,
+            logger=logger,
+            address=address,
+            mode=mode,
+            update_method=update_method,
+            needs_poll_method=needs_poll_method,
+            poll_method=poll_method,
+            poll_debouncer=poll_debouncer,
+            connectable=connectable,
+        )
+        self.discovered_device_classes = discovered_device_classes
+        self.device_data = device_data
+
+
+class XiaomiPassiveBluetoothDataProcessor(PassiveBluetoothDataProcessor):
+    """Define a Xiaomi Bluetooth Passive Update Data Processor."""
+
+    coordinator: XiaomiActiveBluetoothProcessorCoordinator
diff --git a/homeassistant/components/xiaomi_ble/sensor.py b/homeassistant/components/xiaomi_ble/sensor.py
index 84ef91bf5a8..f0f0d7fa71e 100644
--- a/homeassistant/components/xiaomi_ble/sensor.py
+++ b/homeassistant/components/xiaomi_ble/sensor.py
@@ -5,9 +5,7 @@ from xiaomi_ble import DeviceClass, SensorUpdate, Units
 
 from homeassistant import config_entries
 from homeassistant.components.bluetooth.passive_update_processor import (
-    PassiveBluetoothDataProcessor,
     PassiveBluetoothDataUpdate,
-    PassiveBluetoothProcessorCoordinator,
     PassiveBluetoothProcessorEntity,
 )
 from homeassistant.components.sensor import (
@@ -33,6 +31,10 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.sensor import sensor_device_info_to_hass_device_info
 
 from .const import DOMAIN
+from .coordinator import (
+    XiaomiActiveBluetoothProcessorCoordinator,
+    XiaomiPassiveBluetoothDataProcessor,
+)
 from .device import device_key_to_bluetooth_entity_key
 
 SENSOR_DESCRIPTIONS = {
@@ -170,10 +172,12 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Xiaomi BLE sensors."""
-    coordinator: PassiveBluetoothProcessorCoordinator = hass.data[DOMAIN][
+    coordinator: XiaomiActiveBluetoothProcessorCoordinator = hass.data[DOMAIN][
         entry.entry_id
     ]
-    processor = PassiveBluetoothDataProcessor(sensor_update_to_bluetooth_data_update)
+    processor = XiaomiPassiveBluetoothDataProcessor(
+        sensor_update_to_bluetooth_data_update
+    )
     entry.async_on_unload(
         processor.async_add_entities_listener(
             XiaomiBluetoothSensorEntity, async_add_entities
@@ -183,7 +187,7 @@ async def async_setup_entry(
 
 
 class XiaomiBluetoothSensorEntity(
-    PassiveBluetoothProcessorEntity[PassiveBluetoothDataProcessor[float | int | None]],
+    PassiveBluetoothProcessorEntity[XiaomiPassiveBluetoothDataProcessor],
     SensorEntity,
 ):
     """Representation of a xiaomi ble sensor."""
@@ -192,3 +196,11 @@ class XiaomiBluetoothSensorEntity(
     def native_value(self) -> int | float | None:
         """Return the native value."""
         return self.processor.entity_data.get(self.entity_key)
+
+    @property
+    def available(self) -> bool:
+        """Return True if entity is available."""
+        coordinator: XiaomiActiveBluetoothProcessorCoordinator = (
+            self.processor.coordinator
+        )
+        return coordinator.device_data.sleepy_device or super().available
diff --git a/tests/components/xiaomi_ble/test_sensor.py b/tests/components/xiaomi_ble/test_sensor.py
index fff8d9b20f1..7f39228a012 100644
--- a/tests/components/xiaomi_ble/test_sensor.py
+++ b/tests/components/xiaomi_ble/test_sensor.py
@@ -1,8 +1,20 @@
 """Test Xiaomi BLE sensors."""
+from datetime import timedelta
+import time
+from unittest.mock import patch
+
+from homeassistant.components.bluetooth import (
+    FALLBACK_MAXIMUM_STALE_ADVERTISEMENT_SECONDS,
+)
 from homeassistant.components.sensor import ATTR_STATE_CLASS
 from homeassistant.components.xiaomi_ble.const import DOMAIN
-from homeassistant.const import ATTR_FRIENDLY_NAME, ATTR_UNIT_OF_MEASUREMENT
+from homeassistant.const import (
+    ATTR_FRIENDLY_NAME,
+    ATTR_UNIT_OF_MEASUREMENT,
+    STATE_UNAVAILABLE,
+)
 from homeassistant.core import HomeAssistant
+from homeassistant.util import dt as dt_util
 
 from . import (
     HHCCJCY10_SERVICE_INFO,
@@ -12,8 +24,11 @@ from . import (
     make_advertisement,
 )
 
-from tests.common import MockConfigEntry
-from tests.components.bluetooth import inject_bluetooth_service_info_bleak
+from tests.common import MockConfigEntry, async_fire_time_changed
+from tests.components.bluetooth import (
+    inject_bluetooth_service_info_bleak,
+    patch_all_discovered_devices,
+)
 
 
 async def test_sensors(hass: HomeAssistant) -> None:
@@ -610,3 +625,103 @@ async def test_miscale_v2_uuid(hass: HomeAssistant) -> None:
 
     assert await hass.config_entries.async_unload(entry.entry_id)
     await hass.async_block_till_done()
+
+
+async def test_unavailable(hass: HomeAssistant) -> None:
+    """Test normal device goes to unavailable after 60 minutes."""
+    start_monotonic = time.monotonic()
+
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        unique_id="58:2D:34:12:20:89",
+        data={"bindkey": "a3bfe9853dd85a620debe3620caaa351"},
+    )
+    entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_all()) == 0
+    inject_bluetooth_service_info_bleak(
+        hass,
+        make_advertisement(
+            "58:2D:34:12:20:89",
+            b"XXo\x06\x07\x89 \x124-X_\x17m\xd5O\x02\x00\x00/\xa4S\xfa",
+        ),
+    )
+    await hass.async_block_till_done()
+    assert len(hass.states.async_all()) == 1
+
+    temp_sensor = hass.states.get("sensor.temperature_humidity_sensor_2089_temperature")
+    assert temp_sensor.state == "22.6"
+
+    # Fastforward time without BLE advertisements
+    monotonic_now = start_monotonic + FALLBACK_MAXIMUM_STALE_ADVERTISEMENT_SECONDS + 1
+
+    with patch(
+        "homeassistant.components.bluetooth.manager.MONOTONIC_TIME",
+        return_value=monotonic_now,
+    ), patch_all_discovered_devices([]):
+        async_fire_time_changed(
+            hass,
+            dt_util.utcnow()
+            + timedelta(seconds=FALLBACK_MAXIMUM_STALE_ADVERTISEMENT_SECONDS + 1),
+        )
+        await hass.async_block_till_done()
+
+    temp_sensor = hass.states.get("sensor.temperature_humidity_sensor_2089_temperature")
+
+    # Sleepy devices should keep their state over time
+    assert temp_sensor.state == STATE_UNAVAILABLE
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+
+async def test_sleepy_device(hass: HomeAssistant) -> None:
+    """Test normal device goes to unavailable after 60 minutes."""
+    start_monotonic = time.monotonic()
+
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        unique_id="50:FB:19:1B:B5:DC",
+    )
+    entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_all()) == 0
+    inject_bluetooth_service_info_bleak(hass, MISCALE_V1_SERVICE_INFO)
+
+    await hass.async_block_till_done()
+    assert len(hass.states.async_all()) == 2
+
+    mass_non_stabilized_sensor = hass.states.get(
+        "sensor.mi_smart_scale_b5dc_mass_non_stabilized"
+    )
+    assert mass_non_stabilized_sensor.state == "86.55"
+
+    # Fastforward time without BLE advertisements
+    monotonic_now = start_monotonic + FALLBACK_MAXIMUM_STALE_ADVERTISEMENT_SECONDS + 1
+
+    with patch(
+        "homeassistant.components.bluetooth.manager.MONOTONIC_TIME",
+        return_value=monotonic_now,
+    ), patch_all_discovered_devices([]):
+        async_fire_time_changed(
+            hass,
+            dt_util.utcnow()
+            + timedelta(seconds=FALLBACK_MAXIMUM_STALE_ADVERTISEMENT_SECONDS + 1),
+        )
+        await hass.async_block_till_done()
+
+    mass_non_stabilized_sensor = hass.states.get(
+        "sensor.mi_smart_scale_b5dc_mass_non_stabilized"
+    )
+
+    # Sleepy devices should keep their state over time
+    assert mass_non_stabilized_sensor.state == "86.55"
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()