Refactor ESPHome Bluetooth connection logic to prepare for esphome-bleak (#105747)

.coveragerc

@@ -337,7 +337,6 @@ omit =
     homeassistant/components/escea/__init__.py
     homeassistant/components/escea/climate.py
     homeassistant/components/escea/discovery.py
-    homeassistant/components/esphome/bluetooth/*
     homeassistant/components/esphome/manager.py
     homeassistant/components/etherscan/sensor.py
     homeassistant/components/eufy/*

homeassistant/components/esphome/bluetooth/__init__.py

@@ -5,9 +5,9 @@ import asyncio
 from collections.abc import Coroutine
 from functools import partial
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from aioesphomeapi import APIClient, BluetoothProxyFeature
+from aioesphomeapi import APIClient, BluetoothProxyFeature, DeviceInfo
 from bleak_esphome.backend.cache import ESPHomeBluetoothCache
 from bleak_esphome.backend.client import ESPHomeClient, ESPHomeClientData
 from bleak_esphome.backend.device import ESPHomeBluetoothDevice
@@ -17,7 +17,6 @@ from homeassistant.components.bluetooth import (
     HaBluetoothConnector,
     async_register_scanner,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback as hass_callback
 
 from ..entry_data import RuntimeEntryData
@@ -25,20 +24,19 @@ from ..entry_data import RuntimeEntryData
 _LOGGER = logging.getLogger(__name__)
 
 
-@hass_callback
-def _async_can_connect(
-    entry_data: RuntimeEntryData, bluetooth_device: ESPHomeBluetoothDevice, source: str
-) -> bool:
+def _async_can_connect(bluetooth_device: ESPHomeBluetoothDevice, source: str) -> bool:
     """Check if a given source can make another connection."""
-    can_connect = bool(entry_data.available and bluetooth_device.ble_connections_free)
+    can_connect = bool(
+        bluetooth_device.available and bluetooth_device.ble_connections_free
+    )
     _LOGGER.debug(
         (
             "%s [%s]: Checking can connect, available=%s, ble_connections_free=%s"
             " result=%s"
         ),
-        entry_data.name,
+        bluetooth_device.name,
         source,
-        entry_data.available,
+        bluetooth_device.available,
         bluetooth_device.ble_connections_free,
         can_connect,
     )
@@ -54,25 +52,25 @@ def _async_unload(unload_callbacks: list[CALLBACK_TYPE]) -> None:
 
 async def async_connect_scanner(
     hass: HomeAssistant,
-    entry: ConfigEntry,
-    cli: APIClient,
     entry_data: RuntimeEntryData,
+    cli: APIClient,
+    device_info: DeviceInfo,
     cache: ESPHomeBluetoothCache,
 ) -> CALLBACK_TYPE:
     """Connect scanner."""
-    assert entry.unique_id is not None
-    source = str(entry.unique_id)
-    device_info = entry_data.device_info
-    assert device_info is not None
-    feature_flags = device_info.bluetooth_proxy_feature_flags_compat(
-        entry_data.api_version
-    )
+    source = device_info.mac_address
+    name = device_info.name
+    if TYPE_CHECKING:
+        assert cli.api_version is not None
+    feature_flags = device_info.bluetooth_proxy_feature_flags_compat(cli.api_version)
     connectable = bool(feature_flags & BluetoothProxyFeature.ACTIVE_CONNECTIONS)
-    bluetooth_device = ESPHomeBluetoothDevice(entry_data.name, device_info.mac_address)
+    bluetooth_device = ESPHomeBluetoothDevice(
+        name, device_info.mac_address, available=entry_data.available
+    )
     entry_data.bluetooth_device = bluetooth_device
     _LOGGER.debug(
         "%s [%s]: Connecting scanner feature_flags=%s, connectable=%s",
-        entry.title,
+        name,
         source,
         feature_flags,
         connectable,
@@ -82,8 +80,8 @@ async def async_connect_scanner(
         cache=cache,
         client=cli,
         device_info=device_info,
-        api_version=entry_data.api_version,
-        title=entry.title,
+        api_version=cli.api_version,
+        title=name,
         scanner=None,
         disconnect_callbacks=entry_data.disconnect_callbacks,
     )
@@ -92,11 +90,9 @@ async def async_connect_scanner(
         # https://github.com/python/mypy/issues/1484
         client=partial(ESPHomeClient, client_data=client_data),  # type: ignore[arg-type]
         source=source,
-        can_connect=hass_callback(
-            partial(_async_can_connect, entry_data, bluetooth_device, source)
-        ),
+        can_connect=partial(_async_can_connect, bluetooth_device, source),
     )
-    scanner = ESPHomeScanner(source, entry.title, connector, connectable)
+    scanner = ESPHomeScanner(source, name, connector, connectable)
     client_data.scanner = scanner
     coros: list[Coroutine[Any, Any, CALLBACK_TYPE]] = []
     # These calls all return a callback that can be used to unsubscribe

homeassistant/components/esphome/entry_data.py

@@ -420,6 +420,8 @@ class RuntimeEntryData:
         Safe to call multiple times.
         """
         self.available = False
+        if self.bluetooth_device:
+            self.bluetooth_device.available = False
         # Make a copy since calling the disconnect callbacks
         # may also try to discard/remove themselves.
         for disconnect_cb in self.disconnect_callbacks.copy():
@@ -428,3 +430,21 @@ class RuntimeEntryData:
         # to it and make sure all the callbacks can be GC'd.
         self.disconnect_callbacks.clear()
         self.disconnect_callbacks = set()
+
+    @callback
+    def async_on_connect(
+        self, device_info: DeviceInfo, api_version: APIVersion
+    ) -> None:
+        """Call when the entry has been connected."""
+        self.available = True
+        if self.bluetooth_device:
+            self.bluetooth_device.available = True
+
+        self.device_info = device_info
+        self.api_version = api_version
+        # Reset expected disconnect flag on successful reconnect
+        # as it will be flipped to False on unexpected disconnect.
+        #
+        # We use this to determine if a deep sleep device should
+        # be marked as unavailable or not.
+        self.expected_disconnect = True

homeassistant/components/esphome/manager.py

@@ -447,16 +447,10 @@ class ESPHomeManager:
                 entry, data={**entry.data, CONF_DEVICE_NAME: device_info.name}
             )
 
-        entry_data.device_info = device_info
-        assert cli.api_version is not None
-        entry_data.api_version = cli.api_version
-        entry_data.available = True
-        # Reset expected disconnect flag on successful reconnect
-        # as it will be flipped to False on unexpected disconnect.
-        #
-        # We use this to determine if a deep sleep device should
-        # be marked as unavailable or not.
-        entry_data.expected_disconnect = True
+        api_version = cli.api_version
+        assert api_version is not None, "API version must be set"
+        entry_data.async_on_connect(device_info, api_version)
+
         if device_info.name:
             reconnect_logic.name = device_info.name
 
@@ -472,10 +466,10 @@ class ESPHomeManager:
         setup_coros_with_disconnect_callbacks: list[
             Coroutine[Any, Any, CALLBACK_TYPE]
         ] = []
-        if device_info.bluetooth_proxy_feature_flags_compat(cli.api_version):
+        if device_info.bluetooth_proxy_feature_flags_compat(api_version):
             setup_coros_with_disconnect_callbacks.append(
                 async_connect_scanner(
-                    hass, entry, cli, entry_data, self.domain_data.bluetooth_cache
+                    hass, entry_data, cli, device_info, self.domain_data.bluetooth_cache
                 )
             )
 
@@ -507,7 +501,7 @@ class ESPHomeManager:
             entry_data.disconnect_callbacks.add(cancel_callback)
 
         hass.async_create_task(entry_data.async_save_to_store())
-        _async_check_firmware_version(hass, device_info, entry_data.api_version)
+        _async_check_firmware_version(hass, device_info, api_version)
         _async_check_using_api_password(hass, device_info, bool(self.password))
 
     async def on_disconnect(self, expected_disconnect: bool) -> None:

tests/components/esphome/bluetooth/__init__.py

@@ -0,0 +1 @@
+"""Bluetooth tests for ESPHome."""

tests/components/esphome/bluetooth/test_client.py

@@ -32,11 +32,11 @@ async def client_data_fixture(
             mac_address=ESP_MAC_ADDRESS,
             name=ESP_NAME,
             bluetooth_proxy_feature_flags=BluetoothProxyFeature.PASSIVE_SCAN
-            & BluetoothProxyFeature.ACTIVE_CONNECTIONS
-            & BluetoothProxyFeature.REMOTE_CACHING
-            & BluetoothProxyFeature.PAIRING
-            & BluetoothProxyFeature.CACHE_CLEARING
-            & BluetoothProxyFeature.RAW_ADVERTISEMENTS,
+            | BluetoothProxyFeature.ACTIVE_CONNECTIONS
+            | BluetoothProxyFeature.REMOTE_CACHING
+            | BluetoothProxyFeature.PAIRING
+            | BluetoothProxyFeature.CACHE_CLEARING
+            | BluetoothProxyFeature.RAW_ADVERTISEMENTS,
         ),
         api_version=APIVersion(1, 9),
         title=ESP_NAME,

tests/components/esphome/bluetooth/test_init.py

@@ -0,0 +1,46 @@
+"""Test the ESPHome bluetooth integration."""
+
+from homeassistant.components import bluetooth
+from homeassistant.core import HomeAssistant
+
+from ..conftest import MockESPHomeDevice
+
+
+async def test_bluetooth_connect_with_raw_adv(
+    hass: HomeAssistant, mock_bluetooth_entry_with_raw_adv: MockESPHomeDevice
+) -> None:
+    """Test bluetooth connect with raw advertisements."""
+    scanner = bluetooth.async_scanner_by_source(hass, "11:22:33:44:55:aa")
+    assert scanner is not None
+    assert scanner.connectable is True
+    assert scanner.scanning is True
+    assert scanner.connector.can_connect() is False  # no connection slots
+    await mock_bluetooth_entry_with_raw_adv.mock_disconnect(True)
+    await hass.async_block_till_done()
+
+    scanner = bluetooth.async_scanner_by_source(hass, "11:22:33:44:55:aa")
+    assert scanner is None
+    await mock_bluetooth_entry_with_raw_adv.mock_connect()
+    await hass.async_block_till_done()
+    scanner = bluetooth.async_scanner_by_source(hass, "11:22:33:44:55:aa")
+    assert scanner.scanning is True
+
+
+async def test_bluetooth_connect_with_legacy_adv(
+    hass: HomeAssistant, mock_bluetooth_entry_with_legacy_adv: MockESPHomeDevice
+) -> None:
+    """Test bluetooth connect with legacy advertisements."""
+    scanner = bluetooth.async_scanner_by_source(hass, "11:22:33:44:55:aa")
+    assert scanner is not None
+    assert scanner.connectable is True
+    assert scanner.scanning is True
+    assert scanner.connector.can_connect() is False  # no connection slots
+    await mock_bluetooth_entry_with_legacy_adv.mock_disconnect(True)
+    await hass.async_block_till_done()
+
+    scanner = bluetooth.async_scanner_by_source(hass, "11:22:33:44:55:aa")
+    assert scanner is None
+    await mock_bluetooth_entry_with_legacy_adv.mock_connect()
+    await hass.async_block_till_done()
+    scanner = bluetooth.async_scanner_by_source(hass, "11:22:33:44:55:aa")
+    assert scanner.scanning is True

tests/components/esphome/conftest.py

@@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, Mock, patch
 from aioesphomeapi import (
     APIClient,
     APIVersion,
+    BluetoothProxyFeature,
     DeviceInfo,
     EntityInfo,
     EntityState,
@@ -311,6 +312,54 @@ async def mock_voice_assistant_v2_entry(mock_voice_assistant_entry) -> MockConfi
     return await mock_voice_assistant_entry(version=2)
 
 
+@pytest.fixture
+async def mock_bluetooth_entry(
+    hass: HomeAssistant,
+    mock_client: APIClient,
+):
+    """Set up an ESPHome entry with bluetooth."""
+
+    async def _mock_bluetooth_entry(
+        bluetooth_proxy_feature_flags: BluetoothProxyFeature
+    ) -> MockESPHomeDevice:
+        return await _mock_generic_device_entry(
+            hass,
+            mock_client,
+            {"bluetooth_proxy_feature_flags": bluetooth_proxy_feature_flags},
+            ([], []),
+            [],
+        )
+
+    return _mock_bluetooth_entry
+
+
+@pytest.fixture
+async def mock_bluetooth_entry_with_raw_adv(mock_bluetooth_entry) -> MockESPHomeDevice:
+    """Set up an ESPHome entry with bluetooth and raw advertisements."""
+    return await mock_bluetooth_entry(
+        bluetooth_proxy_feature_flags=BluetoothProxyFeature.PASSIVE_SCAN
+        | BluetoothProxyFeature.ACTIVE_CONNECTIONS
+        | BluetoothProxyFeature.REMOTE_CACHING
+        | BluetoothProxyFeature.PAIRING
+        | BluetoothProxyFeature.CACHE_CLEARING
+        | BluetoothProxyFeature.RAW_ADVERTISEMENTS
+    )
+
+
+@pytest.fixture
+async def mock_bluetooth_entry_with_legacy_adv(
+    mock_bluetooth_entry
+) -> MockESPHomeDevice:
+    """Set up an ESPHome entry with bluetooth with legacy advertisements."""
+    return await mock_bluetooth_entry(
+        bluetooth_proxy_feature_flags=BluetoothProxyFeature.PASSIVE_SCAN
+        | BluetoothProxyFeature.ACTIVE_CONNECTIONS
+        | BluetoothProxyFeature.REMOTE_CACHING
+        | BluetoothProxyFeature.PAIRING
+        | BluetoothProxyFeature.CACHE_CLEARING
+    )
+
+
 @pytest.fixture
 async def mock_generic_device_entry(
     hass: HomeAssistant,