Fix Yale Access Bluetooth key discovery timing issues (#151433)

Co-authored-by: Copilot <[email protected]>

Author: bdraco

homeassistant/components/yalexs_ble/__init__.py

@@ -19,6 +19,7 @@
 from homeassistant.core import CALLBACK_TYPE, CoreState, Event, HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 
+from .config_cache import async_get_validated_config
 from .const import (
     CONF_ALWAYS_CONNECTED,
     CONF_KEY,
@@ -96,13 +97,30 @@ def _async_shutdown(event: Event | None = None) -> None:
     )
 
     try:
-        await push_lock.wait_for_first_update(DEVICE_TIMEOUT)
-    except AuthError as ex:
-        raise ConfigEntryAuthFailed(str(ex)) from ex
-    except (YaleXSBLEError, TimeoutError) as ex:
-        raise ConfigEntryNotReady(
-            f"{ex}; Try moving the Bluetooth adapter closer to {local_name}"
-        ) from ex
+        await _async_wait_for_first_update(push_lock, local_name)
+    except ConfigEntryAuthFailed:
+        # If key has rotated, try to fetch it from the cache
+        # and update
+        if (validated_config := async_get_validated_config(hass, address)) and (
+            validated_config.key != entry.data[CONF_KEY]
+            or validated_config.slot != entry.data[CONF_SLOT]
+        ):
+            assert shutdown_callback is not None
+            shutdown_callback()
+            push_lock.set_lock_key(validated_config.key, validated_config.slot)
+            shutdown_callback = await push_lock.start()
+            await _async_wait_for_first_update(push_lock, local_name)
+            # If we can use the cached key and slot, update the entry.
+            hass.config_entries.async_update_entry(
+                entry,
+                data={
+                    **entry.data,
+                    CONF_KEY: validated_config.key,
+                    CONF_SLOT: validated_config.slot,
+                },
+            )
+        else:
+            raise
 
     entry.runtime_data = YaleXSBLEData(entry.title, push_lock, always_connected)
 
@@ -135,6 +153,18 @@ def _async_state_changed(
     return True
 
 
+async def _async_wait_for_first_update(push_lock: PushLock, local_name: str) -> None:
+    """Wait for the first update from the push lock."""
+    try:
+        await push_lock.wait_for_first_update(DEVICE_TIMEOUT)
+    except AuthError as ex:
+        raise ConfigEntryAuthFailed(str(ex)) from ex
+    except (YaleXSBLEError, TimeoutError) as ex:
+        raise ConfigEntryNotReady(
+            f"{ex}; Try moving the Bluetooth adapter closer to {local_name}"
+        ) from ex
+
+
 async def async_unload_entry(hass: HomeAssistant, entry: YALEXSBLEConfigEntry) -> bool:
     """Unload a config entry."""
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/yalexs_ble/config_cache.py

@@ -0,0 +1,31 @@
+"""The Yale Access Bluetooth integration."""
+
+from __future__ import annotations
+
+from yalexs_ble import ValidatedLockConfig
+
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.util.hass_dict import HassKey
+
+CONFIG_CACHE: HassKey[dict[str, ValidatedLockConfig]] = HassKey(
+    "yalexs_ble_config_cache"
+)
+
+
+@callback
+def async_add_validated_config(
+    hass: HomeAssistant,
+    address: str,
+    config: ValidatedLockConfig,
+) -> None:
+    """Add a validated config."""
+    hass.data.setdefault(CONFIG_CACHE, {})[address] = config
+
+
+@callback
+def async_get_validated_config(
+    hass: HomeAssistant,
+    address: str,
+) -> ValidatedLockConfig | None:
+    """Get the config for a specific address."""
+    return hass.data.get(CONFIG_CACHE, {}).get(address)

homeassistant/components/yalexs_ble/config_flow.py

@@ -33,6 +33,7 @@
 from homeassistant.data_entry_flow import AbortFlow
 from homeassistant.helpers.typing import DiscoveryInfoType
 
+from .config_cache import async_add_validated_config, async_get_validated_config
 from .const import CONF_ALWAYS_CONNECTED, CONF_KEY, CONF_LOCAL_NAME, CONF_SLOT, DOMAIN
 from .util import async_find_existing_service_info, human_readable_name
 
@@ -92,7 +93,10 @@ async def async_step_bluetooth(
                 None, discovery_info.name, discovery_info.address
             ),
         }
-        return await self.async_step_user()
+        if lock_cfg := async_get_validated_config(self.hass, discovery_info.address):
+            self._lock_cfg = lock_cfg
+            return await self.async_step_integration_discovery_confirm()
+        return await self.async_step_key_slot()
 
     async def async_step_integration_discovery(
         self, discovery_info: DiscoveryInfoType
@@ -105,6 +109,7 @@ async def async_step_integration_discovery(
             discovery_info["key"],
             discovery_info["slot"],
         )
+        async_add_validated_config(self.hass, lock_cfg.address, lock_cfg)
 
         address = lock_cfg.address
         self.local_name = lock_cfg.local_name
@@ -232,30 +237,36 @@ async def async_step_reauth_validate(
             errors=errors,
         )
 
-    async def async_step_user(
+    async def async_step_key_slot(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
-        """Handle the user step to pick discovered device."""
+        """Handle the key and slot step."""
         errors: dict[str, str] = {}
+        discovery_info = self._discovery_info
+        assert discovery_info is not None
+        address = discovery_info.address
+        validated_config = async_get_validated_config(self.hass, address)
 
-        if user_input is not None:
-            self.active = True
-            address = user_input[CONF_ADDRESS]
-            discovery_info = self._discovered_devices[address]
+        if user_input is not None or validated_config:
             local_name = discovery_info.name
-            key = user_input[CONF_KEY]
-            slot = user_input[CONF_SLOT]
-            await self.async_set_unique_id(
-                discovery_info.address, raise_on_progress=False
-            )
+            if validated_config:
+                key = validated_config.key
+                slot = validated_config.slot
+                title = validated_config.name
+            else:
+                assert user_input is not None
+                key = user_input[CONF_KEY]
+                slot = user_input[CONF_SLOT]
+                title = human_readable_name(None, local_name, address)
+            await self.async_set_unique_id(address, raise_on_progress=False)
             self._abort_if_unique_id_configured()
             if not (
                 errors := await async_validate_lock_or_error(
                     local_name, discovery_info.device, key, slot
                 )
             ):
                 return self.async_create_entry(
-                    title=local_name,
+                    title=title,
                     data={
                         CONF_LOCAL_NAME: discovery_info.name,
                         CONF_ADDRESS: discovery_info.address,
@@ -264,24 +275,48 @@ async def async_step_user(
                     },
                 )
 
-        if discovery := self._discovery_info:
+        return self.async_show_form(
+            step_id="key_slot",
+            data_schema=vol.Schema(
+                {
+                    vol.Required(CONF_KEY): str,
+                    vol.Required(CONF_SLOT): int,
+                }
+            ),
+            errors=errors,
+            description_placeholders={
+                "address": address,
+                "title": self._async_get_name_from_address(address),
+            },
+        )
+
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Handle the user step to pick discovered device."""
+        errors: dict[str, str] = {}
+
+        if user_input is not None:
+            self.active = True
+            address = user_input[CONF_ADDRESS]
+            self._discovery_info = self._discovered_devices[address]
+            return await self.async_step_key_slot()
+
+        current_addresses = self._async_current_ids(include_ignore=False)
+        current_unique_names = {
+            entry.data.get(CONF_LOCAL_NAME)
+            for entry in self._async_current_entries()
+            if local_name_is_unique(entry.data.get(CONF_LOCAL_NAME))
+        }
+        for discovery in async_discovered_service_info(self.hass):
+            if (
+                discovery.address in current_addresses
+                or discovery.name in current_unique_names
+                or discovery.address in self._discovered_devices
+                or YALE_MFR_ID not in discovery.manufacturer_data
+            ):
+                continue
             self._discovered_devices[discovery.address] = discovery
-        else:
-            current_addresses = self._async_current_ids(include_ignore=False)
-            current_unique_names = {
-                entry.data.get(CONF_LOCAL_NAME)
-                for entry in self._async_current_entries()
-                if local_name_is_unique(entry.data.get(CONF_LOCAL_NAME))
-            }
-            for discovery in async_discovered_service_info(self.hass):
-                if (
-                    discovery.address in current_addresses
-                    or discovery.name in current_unique_names
-                    or discovery.address in self._discovered_devices
-                    or YALE_MFR_ID not in discovery.manufacturer_data
-                ):
-                    continue
-                self._discovered_devices[discovery.address] = discovery
 
         if not self._discovered_devices:
             return self.async_abort(reason="no_devices_found")
@@ -290,14 +325,12 @@ async def async_step_user(
             {
                 vol.Required(CONF_ADDRESS): vol.In(
                     {
-                        service_info.address: (
-                            f"{service_info.name} ({service_info.address})"
+                        service_info.address: self._async_get_name_from_address(
+                            service_info.address
                         )
                         for service_info in self._discovered_devices.values()
                     }
-                ),
-                vol.Required(CONF_KEY): str,
-                vol.Required(CONF_SLOT): int,
+                )
             }
         )
         return self.async_show_form(
@@ -306,6 +339,18 @@ async def async_step_user(
             errors=errors,
         )
 
+    @callback
+    def _async_get_name_from_address(self, address: str) -> str:
+        """Get the name of a device from its address."""
+        if validated_config := async_get_validated_config(self.hass, address):
+            return f"{validated_config.name} ({address})"
+        if address in self._discovered_devices:
+            service_info = self._discovered_devices[address]
+            return f"{service_info.name} ({service_info.address})"
+        assert self._discovery_info is not None
+        assert self._discovery_info.address == address
+        return f"{self._discovery_info.name} ({address})"
+
     @staticmethod
     @callback
     def async_get_options_flow(

homeassistant/components/yalexs_ble/strings.json

@@ -3,18 +3,23 @@
     "flow_title": "{name}",
     "step": {
       "user": {
-        "description": "Check the documentation for how to find the offline key. If you are using the August cloud integration to obtain the key, you may need to reload the August cloud integration while the lock is in Bluetooth range.",
+        "description": "Select the device you want to set up over Bluetooth.",
+        "data": {
+          "address": "Bluetooth address"
+        }
+      },
+      "key_slot": {
+        "description": "Enter the key for the {title} lock with address {address}. If you are using the August or Yale cloud integration to obtain the key, you may be able to avoid this manual setup by reloading the August or Yale cloud integration while the lock is in Bluetooth range.",
         "data": {
-          "address": "Bluetooth address",
           "key": "Offline Key (32-byte hex string)",
           "slot": "Offline Key Slot (Integer between 0 and 255)"
         }
       },
       "reauth_validate": {
-        "description": "Enter the updated key for the {title} lock with address {address}. If you are using the August cloud integration to obtain the key, you may be able to avoid manual reauthentication by reloading the August cloud integration while the lock is in Bluetooth range.",
+        "description": "Enter the updated key for the {title} lock with address {address}. If you are using the August or Yale cloud integration to obtain the key, you may be able to avoid manual re-authentication by reloading the August or Yale cloud integration while the lock is in Bluetooth range.",
         "data": {
-          "key": "[%key:component::yalexs_ble::config::step::user::data::key%]",
-          "slot": "[%key:component::yalexs_ble::config::step::user::data::slot%]"
+          "key": "[%key:component::yalexs_ble::config::step::key_slot::data::key%]",
+          "slot": "[%key:component::yalexs_ble::config::step::key_slot::data::slot%]"
         }
       },
       "integration_discovery_confirm": {