Update oralb to show battery percentage (#85800)

Co-authored-by: J. Nick Koston <nick@koston.org>
fixes undefined

CODEOWNERS

@@ -851,8 +851,8 @@ build.json @home-assistant/supervisor
 /tests/components/openweathermap/ @fabaff @freekode @nzapponi
 /homeassistant/components/opnsense/ @mtreinish
 /tests/components/opnsense/ @mtreinish
-/homeassistant/components/oralb/ @bdraco
-/tests/components/oralb/ @bdraco
+/homeassistant/components/oralb/ @bdraco @conway20
+/tests/components/oralb/ @bdraco @conway20
 /homeassistant/components/oru/ @bvlaicu
 /homeassistant/components/overkiz/ @imicknl @vlebourl @tetienne @nyroDev
 /tests/components/overkiz/ @imicknl @vlebourl @tetienne @nyroDev

homeassistant/components/oralb/__init__.py

@@ -5,13 +5,17 @@ import logging
 
 from oralb_ble import OralBBluetoothDeviceData
 
-from homeassistant.components.bluetooth import BluetoothScanningMode
-from homeassistant.components.bluetooth.passive_update_processor import (
-    PassiveBluetoothProcessorCoordinator,
+from homeassistant.components.bluetooth import (
+    BluetoothScanningMode,
+    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 HomeAssistant
+from homeassistant.core import CoreState, HomeAssistant
 
 from .const import DOMAIN
 
@@ -25,14 +29,55 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     address = entry.unique_id
     assert address is not None
     data = OralBBluetoothDeviceData()
+
+    def _needs_poll(
+        service_info: BluetoothServiceInfoBleak, last_poll: float | None
+    ) -> bool:
+        # Only poll if hass is running, we need to poll,
+        # and we actually have a way to connect to the device
+        return (
+            hass.state == CoreState.running
+            and data.poll_needed(service_info, last_poll)
+            and bool(
+                async_ble_device_from_address(
+                    hass, service_info.device.address, connectable=True
+                )
+            )
+        )
+
+    async def _async_poll(service_info: BluetoothServiceInfoBleak):
+        # BluetoothServiceInfoBleak is defined in HA, otherwise would just pass it
+        # directly to the Xiaomi code
+        # Make sure the device we have is one that we can connect with
+        # in case its coming from a passive scanner
+        if service_info.connectable:
+            connectable_device = service_info.device
+        elif device := async_ble_device_from_address(
+            hass, service_info.device.address, True
+        ):
+            connectable_device = device
+        else:
+            # We have no bluetooth controller that is in range of
+            # the device to poll it
+            raise RuntimeError(
+                f"No connectable device found for {service_info.device.address}"
+            )
+        return await data.async_poll(connectable_device)
+
     coordinator = hass.data.setdefault(DOMAIN, {})[
         entry.entry_id
-    ] = PassiveBluetoothProcessorCoordinator(
+    ] = ActiveBluetoothProcessorCoordinator(
         hass,
         _LOGGER,
         address=address,
         mode=BluetoothScanningMode.PASSIVE,
         update_method=data.update,
+        needs_poll_method=_needs_poll,
+        poll_method=_async_poll,
+        # We will take advertisements from non-connectable devices
+        # since we will trade the BLEDevice for a connectable one
+        # if we need to poll it
+        connectable=False,
     )
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     entry.async_on_unload(

homeassistant/components/oralb/manifest.json

@@ -8,8 +8,8 @@
       "manufacturer_id": 220
     }
   ],
-  "requirements": ["oralb-ble==0.14.3"],
+  "requirements": ["oralb-ble==0.17.1"],
   "dependencies": ["bluetooth"],
-  "codeowners": ["@bdraco"],
+  "codeowners": ["@bdraco", "@conway20"],
   "iot_class": "local_push"
 }

homeassistant/components/oralb/sensor.py

@@ -18,7 +18,11 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.const import SIGNAL_STRENGTH_DECIBELS_MILLIWATT, UnitOfTime
+from homeassistant.const import (
+    PERCENTAGE,
+    SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
+    UnitOfTime,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -59,6 +63,12 @@ SENSOR_DESCRIPTIONS: dict[str, SensorEntityDescription] = {
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
     ),
+    OralBSensor.BATTERY_PERCENT: SensorEntityDescription(
+        key=OralBSensor.BATTERY_PERCENT,
+        device_class=SensorDeviceClass.BATTERY,
+        native_unit_of_measurement=PERCENTAGE,
+        state_class=SensorStateClass.MEASUREMENT,
+    ),
 }
 
 

requirements_all.txt

@@ -1290,7 +1290,7 @@ openwrt-luci-rpc==1.1.11
 openwrt-ubus-rpc==0.0.2
 
 # homeassistant.components.oralb
-oralb-ble==0.14.3
+oralb-ble==0.17.1
 
 # homeassistant.components.oru
 oru==0.1.11

requirements_test_all.txt

@@ -938,7 +938,7 @@ open-meteo==0.2.1
 openerz-api==0.1.0
 
 # homeassistant.components.oralb
-oralb-ble==0.14.3
+oralb-ble==0.17.1
 
 # homeassistant.components.ovo_energy
 ovoenergy==1.2.0

tests/components/oralb/__init__.py

@@ -1,8 +1,12 @@
 """Tests for the OralB integration."""
 
+from bleak.backends.device import BLEDevice
+from home_assistant_bluetooth import BluetoothServiceInfoBleak
 
 from homeassistant.helpers.service_info.bluetooth import BluetoothServiceInfo
 
+from tests.components.bluetooth import generate_advertisement_data
+
 NOT_ORALB_SERVICE_INFO = BluetoothServiceInfo(
     name="Not it",
     address="61DE521B-F0BF-9F44-64D4-75BBE1738105",
@@ -33,3 +37,17 @@ ORALB_IO_SERIES_4_SERVICE_INFO = BluetoothServiceInfo(
     service_data={},
     source="local",
 )
+
+ORALB_IO_SERIES_6_SERVICE_INFO = BluetoothServiceInfoBleak(
+    name="Oral-B Toothbrush",
+    address="B0:D2:78:20:1D:CF",
+    device=BLEDevice("B0:D2:78:20:1D:CF", "Oral-B Toothbrush"),
+    rssi=-56,
+    manufacturer_data={220: b"\x062k\x02r\x00\x00\x02\x01\x00\x04"},
+    service_data={"a0f0ff00-5047-4d53-8208-4f72616c2d42": bytearray(b"1\x00\x00\x00")},
+    service_uuids=["a0f0ff00-5047-4d53-8208-4f72616c2d42"],
+    source="local",
+    advertisement=generate_advertisement_data(local_name="Not it"),
+    time=0,
+    connectable=True,
+)

tests/components/oralb/conftest.py

@@ -1,8 +1,53 @@
 """OralB session fixtures."""
 
+from unittest import mock
+
 import pytest
 
 
+class MockServices:
+    """Mock GATTServicesCollection."""
+
+    def get_characteristic(self, key: str) -> str:
+        """Mock GATTServicesCollection.get_characteristic."""
+        return key
+
+
+class MockBleakClient:
+    """Mock BleakClient."""
+
+    services = MockServices()
+
+    def __init__(self, *args, **kwargs):
+        """Mock BleakClient."""
+
+    async def __aenter__(self, *args, **kwargs):
+        """Mock BleakClient.__aenter__."""
+        return self
+
+    async def __aexit__(self, *args, **kwargs):
+        """Mock BleakClient.__aexit__."""
+
+    async def connect(self, *args, **kwargs):
+        """Mock BleakClient.connect."""
+
+    async def disconnect(self, *args, **kwargs):
+        """Mock BleakClient.disconnect."""
+
+
+class MockBleakClientBattery49(MockBleakClient):
+    """Mock BleakClient that returns a battery level of 49."""
+
+    async def read_gatt_char(self, *args, **kwargs) -> bytes:
+        """Mock BleakClient.read_gatt_char."""
+        return b"\x31\x00"
+
+
 @pytest.fixture(autouse=True)
 def mock_bluetooth(enable_bluetooth):
     """Auto mock bluetooth."""
+
+    with mock.patch(
+        "oralb_ble.parser.BleakClientWithServiceCache", MockBleakClientBattery49
+    ):
+        yield

tests/components/oralb/test_sensor.py

@@ -4,10 +4,17 @@
 from homeassistant.components.oralb.const import DOMAIN
 from homeassistant.const import ATTR_FRIENDLY_NAME
 
-from . import ORALB_IO_SERIES_4_SERVICE_INFO, ORALB_SERVICE_INFO
+from . import (
+    ORALB_IO_SERIES_4_SERVICE_INFO,
+    ORALB_IO_SERIES_6_SERVICE_INFO,
+    ORALB_SERVICE_INFO,
+)
 
 from tests.common import MockConfigEntry
-from tests.components.bluetooth import inject_bluetooth_service_info
+from tests.components.bluetooth import (
+    inject_bluetooth_service_info,
+    inject_bluetooth_service_info_bleak,
+)
 
 
 async def test_sensors(hass, entity_registry_enabled_by_default):
@@ -24,7 +31,7 @@ async def test_sensors(hass, entity_registry_enabled_by_default):
     assert len(hass.states.async_all("sensor")) == 0
     inject_bluetooth_service_info(hass, ORALB_SERVICE_INFO)
     await hass.async_block_till_done()
-    assert len(hass.states.async_all("sensor")) == 8
+    assert len(hass.states.async_all("sensor")) == 9
 
     toothbrush_sensor = hass.states.get(
         "sensor.smart_series_7000_48be_toothbrush_state"
@@ -54,7 +61,7 @@ async def test_sensors_io_series_4(hass, entity_registry_enabled_by_default):
     assert len(hass.states.async_all("sensor")) == 0
     inject_bluetooth_service_info(hass, ORALB_IO_SERIES_4_SERVICE_INFO)
     await hass.async_block_till_done()
-    assert len(hass.states.async_all("sensor")) == 8
+    assert len(hass.states.async_all("sensor")) == 9
 
     toothbrush_sensor = hass.states.get("sensor.io_series_4_48be_mode")
     toothbrush_sensor_attrs = toothbrush_sensor.attributes
@@ -63,3 +70,24 @@ async def test_sensors_io_series_4(hass, entity_registry_enabled_by_default):
 
     assert await hass.config_entries.async_unload(entry.entry_id)
     await hass.async_block_till_done()
+
+
+async def test_sensors_battery(hass):
+    """Test receiving battery percentage."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        unique_id=ORALB_IO_SERIES_6_SERVICE_INFO.address,
+    )
+    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, ORALB_IO_SERIES_6_SERVICE_INFO)
+    await hass.async_block_till_done()
+    assert len(hass.states.async_all()) == 7
+
+    bat_sensor = hass.states.get("sensor.io_series_6_7_1dcf_battery")
+    assert bat_sensor.state == "49"
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()