Abort Improv via BLE bluetooth flow if device is provisioned (#102656)

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

homeassistant/components/improv_ble/config_flow.py

@@ -20,13 +20,10 @@ from improv_ble_client import (
 import voluptuous as vol
 
 from homeassistant import config_entries
-from homeassistant.components.bluetooth import (
+from homeassistant.components import bluetooth
-    BluetoothServiceInfoBleak,
-    async_discovered_service_info,
-    async_last_service_info,
-)
 from homeassistant.const import CONF_ADDRESS
-from homeassistant.data_entry_flow import FlowResult
+from homeassistant.core import callback
+from homeassistant.data_entry_flow import AbortFlow, FlowResult
 
 from .const import DOMAIN
 
@@ -42,14 +39,6 @@ STEP_PROVISION_SCHEMA = vol.Schema(
 )
 
 
-class AbortFlow(Exception):
-    """Raised when a flow should be aborted."""
-
-    def __init__(self, reason: str) -> None:
-        """Initialize."""
-        self.reason = reason
-
-
 @dataclass
 class Credentials:
     """Container for WiFi credentials."""
@@ -69,15 +58,16 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     _provision_result: FlowResult | None = None
     _provision_task: asyncio.Task | None = None
     _reauth_entry: config_entries.ConfigEntry | None = None
+    _remove_bluetooth_callback: Callable[[], None] | None = None
     _unsub: Callable[[], None] | None = None
 
     def __init__(self) -> None:
         """Initialize the config flow."""
         self._device: ImprovBLEClient | None = None
         # Populated by user step
-        self._discovered_devices: dict[str, BluetoothServiceInfoBleak] = {}
+        self._discovered_devices: dict[str, bluetooth.BluetoothServiceInfoBleak] = {}
         # Populated by bluetooth, reauth_confirm and user steps
-        self._discovery_info: BluetoothServiceInfoBleak | None = None
+        self._discovery_info: bluetooth.BluetoothServiceInfoBleak | None = None
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
@@ -95,7 +85,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             return await self.async_step_start_improv()
 
         current_addresses = self._async_current_ids()
-        for discovery in async_discovered_service_info(self.hass):
+        for discovery in bluetooth.async_discovered_service_info(self.hass):
             if (
                 discovery.address in current_addresses
                 or discovery.address in self._discovered_devices
@@ -125,22 +115,66 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             errors=errors,
         )
 
-    async def async_step_bluetooth(
+    def _abort_if_provisioned(self) -> None:
-        self, discovery_info: BluetoothServiceInfoBleak
+        """Check improv state and abort flow if needed."""
-    ) -> FlowResult:
+        # mypy is not aware that we can't get here without having these set already
-        """Handle the Bluetooth discovery step."""
+        assert self._discovery_info is not None
-        await self.async_set_unique_id(discovery_info.address)
+
-        self._abort_if_unique_id_configured()
+        service_data = self._discovery_info.service_data
-        service_data = discovery_info.service_data
         improv_service_data = ImprovServiceData.from_bytes(
             service_data[SERVICE_DATA_UUID]
         )
         if improv_service_data.state in (State.PROVISIONING, State.PROVISIONED):
             _LOGGER.debug(
-                "Device is already provisioned: %s", improv_service_data.state
+                "Aborting improv flow, device is already provisioned: %s",
+                improv_service_data.state,
             )
-            return self.async_abort(reason="already_provisioned")
+            raise AbortFlow("already_provisioned")
+
+    @callback
+    def _async_update_ble(
+        self,
+        service_info: bluetooth.BluetoothServiceInfoBleak,
+        change: bluetooth.BluetoothChange,
+    ) -> None:
+        """Update from a ble callback."""
+        _LOGGER.debug(
+            "Got updated BLE data: %s",
+            service_info.service_data[SERVICE_DATA_UUID].hex(),
+        )
+
+        self._discovery_info = service_info
+        try:
+            self._abort_if_provisioned()
+        except AbortFlow:
+            self.hass.config_entries.flow.async_abort(self.flow_id)
+
+    def _unregister_bluetooth_callback(self) -> None:
+        """Unregister bluetooth callbacks."""
+        if not self._remove_bluetooth_callback:
+            return
+        self._remove_bluetooth_callback()
+        self._remove_bluetooth_callback = None
+
+    async def async_step_bluetooth(
+        self, discovery_info: bluetooth.BluetoothServiceInfoBleak
+    ) -> FlowResult:
+        """Handle the Bluetooth discovery step."""
         self._discovery_info = discovery_info
+
+        await self.async_set_unique_id(discovery_info.address)
+        self._abort_if_unique_id_configured()
+        self._abort_if_provisioned()
+
+        self._remove_bluetooth_callback = bluetooth.async_register_callback(
+            self.hass,
+            self._async_update_ble,
+            bluetooth.BluetoothCallbackMatcher(
+                {bluetooth.match.ADDRESS: discovery_info.address}
+            ),
+            bluetooth.BluetoothScanningMode.PASSIVE,
+        )
+
         name = self._discovery_info.name or self._discovery_info.address
         self.context["title_placeholders"] = {"name": name}
         return await self.async_step_bluetooth_confirm()
@@ -159,6 +193,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 description_placeholders={"name": name},
             )
 
+        self._unregister_bluetooth_callback()
         return await self.async_step_start_improv()
 
     async def async_step_start_improv(
@@ -171,23 +206,9 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """
         # mypy is not aware that we can't get here without having these set already
         assert self._discovery_info is not None
-        discovery_info = self._discovery_info = async_last_service_info(
-            self.hass, self._discovery_info.address
-        )
-        if not discovery_info:
-            return self.async_abort(reason="cannot_connect")
-        service_data = discovery_info.service_data
-        improv_service_data = ImprovServiceData.from_bytes(
-            service_data[SERVICE_DATA_UUID]
-        )
-        if improv_service_data.state in (State.PROVISIONING, State.PROVISIONED):
-            _LOGGER.debug(
-                "Device is already provisioned: %s", improv_service_data.state
-            )
-            return self.async_abort(reason="already_provisioned")
 
         if not self._device:
-            self._device = ImprovBLEClient(discovery_info.device)
+            self._device = ImprovBLEClient(self._discovery_info.device)
         device = self._device
 
         if self._can_identify is None:
@@ -387,3 +408,8 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         except Exception as err:  # pylint: disable=broad-except
             _LOGGER.exception("Unexpected exception")
             raise AbortFlow("unknown") from err
+
+    @callback
+    def async_remove(self) -> None:
+        """Notification that the flow has been removed."""
+        self._unregister_bluetooth_callback()

homeassistant/components/improv_ble/strings.json

@@ -38,7 +38,6 @@
     },
     "abort": {
       "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
-      "already_provisioned": "The device is already connected to Wi-Fi. If you want to connect it to another network, try factory resetting it first.",
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
       "characteristic_missing": "The device is either already connected to Wi-Fi, or no longer able to connect to Wi-Fi. If you want to connect it to another network, try factory resetting it first.",
       "no_devices_found": "[%key:common::config_flow::abort::no_devices_found%]",

tests/components/improv_ble/test_config_flow.py

@@ -7,6 +7,7 @@ from improv_ble_client import Error, State, errors as improv_ble_errors
 import pytest
 
 from homeassistant import config_entries
+from homeassistant.components.bluetooth import BluetoothChange
 from homeassistant.components.improv_ble.const import DOMAIN
 from homeassistant.const import CONF_ADDRESS
 from homeassistant.core import HomeAssistant
@@ -36,7 +37,7 @@ async def test_user_step_success(
 ) -> None:
     """Test user step success path."""
     with patch(
-        f"{IMPROV_BLE}.config_flow.async_discovered_service_info",
+        f"{IMPROV_BLE}.config_flow.bluetooth.async_discovered_service_info",
         return_value=[NOT_IMPROV_BLE_DISCOVERY_INFO, IMPROV_BLE_DISCOVERY_INFO],
     ):
         result = await hass.config_entries.flow.async_init(
@@ -46,24 +47,15 @@ async def test_user_step_success(
     assert result["step_id"] == "user"
     assert result["errors"] == {}
 
-    with patch(
+    await _test_common_success_wo_identify(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
+        hass, result, IMPROV_BLE_DISCOVERY_INFO.address, url, abort_reason, placeholders
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
+    )
-    ):
-        await _test_common_success_wo_identify(
-            hass,
-            result,
-            IMPROV_BLE_DISCOVERY_INFO.address,
-            url,
-            abort_reason,
-            placeholders,
-        )
 
 
 async def test_user_step_success_authorize(hass: HomeAssistant) -> None:
     """Test user step success path."""
     with patch(
-        f"{IMPROV_BLE}.config_flow.async_discovered_service_info",
+        f"{IMPROV_BLE}.config_flow.bluetooth.async_discovered_service_info",
         return_value=[NOT_IMPROV_BLE_DISCOVERY_INFO, IMPROV_BLE_DISCOVERY_INFO],
     ):
         result = await hass.config_entries.flow.async_init(
@@ -73,19 +65,15 @@ async def test_user_step_success_authorize(hass: HomeAssistant) -> None:
     assert result["step_id"] == "user"
     assert result["errors"] == {}
 
-    with patch(
+    await _test_common_success_wo_identify_w_authorize(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
+        hass, result, IMPROV_BLE_DISCOVERY_INFO.address
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
+    )
-    ):
-        await _test_common_success_wo_identify_w_authorize(
-            hass, result, IMPROV_BLE_DISCOVERY_INFO.address
-        )
 
 
 async def test_user_step_no_devices_found(hass: HomeAssistant) -> None:
     """Test user step with no devices found."""
     with patch(
-        f"{IMPROV_BLE}.config_flow.async_discovered_service_info",
+        f"{IMPROV_BLE}.config_flow.bluetooth.async_discovered_service_info",
         return_value=[
             PROVISIONED_IMPROV_BLE_DISCOVERY_INFO,
             NOT_IMPROV_BLE_DISCOVERY_INFO,
@@ -111,7 +99,7 @@ async def test_async_step_user_takes_precedence_over_discovery(
     assert result["step_id"] == "bluetooth_confirm"
 
     with patch(
-        f"{IMPROV_BLE}.config_flow.async_discovered_service_info",
+        f"{IMPROV_BLE}.config_flow.bluetooth.async_discovered_service_info",
         return_value=[IMPROV_BLE_DISCOVERY_INFO],
     ):
         result = await hass.config_entries.flow.async_init(
@@ -120,13 +108,9 @@ async def test_async_step_user_takes_precedence_over_discovery(
         )
         assert result["type"] == FlowResultType.FORM
 
-    with patch(
+    await _test_common_success_wo_identify(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
+        hass, result, IMPROV_BLE_DISCOVERY_INFO.address
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
+    )
-    ):
-        await _test_common_success_wo_identify(
-            hass, result, IMPROV_BLE_DISCOVERY_INFO.address
-        )
 
     # Verify the discovery flow was aborted
     assert not hass.config_entries.flow.async_progress(DOMAIN)
@@ -143,48 +127,25 @@ async def test_bluetooth_step_provisioned_device(hass: HomeAssistant) -> None:
     assert result["reason"] == "already_provisioned"
 
 
-async def test_bluetooth_confirm_provisioned_device(hass: HomeAssistant) -> None:
+async def test_bluetooth_step_provisioned_device_2(hass: HomeAssistant) -> None:
-    """Test bluetooth confirm step when device is already provisioned."""
+    """Test bluetooth step when device changes to provisioned."""
-    result = await hass.config_entries.flow.async_init(
+    with patch(
-        DOMAIN,
+        f"{IMPROV_BLE}.config_flow.bluetooth.async_register_callback",
-        context={"source": config_entries.SOURCE_BLUETOOTH},
+    ) as mock_async_register_callback:
-        data=IMPROV_BLE_DISCOVERY_INFO,
+        result = await hass.config_entries.flow.async_init(
-    )
+            DOMAIN,
+            context={"source": config_entries.SOURCE_BLUETOOTH},
+            data=IMPROV_BLE_DISCOVERY_INFO,
+        )
     assert result["type"] == FlowResultType.FORM
     assert result["step_id"] == "bluetooth_confirm"
-    assert result["errors"] is None
 
-    with patch(
+    assert len(hass.config_entries.flow.async_progress_by_handler("improv_ble")) == 1
-        f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", return_value=False
-    ), patch(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
-        return_value=PROVISIONED_IMPROV_BLE_DISCOVERY_INFO,
-    ):
-        result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
-        assert result["type"] == FlowResultType.ABORT
-        assert result["reason"] == "already_provisioned"
 
+    callback = mock_async_register_callback.call_args.args[1]
+    callback(PROVISIONED_IMPROV_BLE_DISCOVERY_INFO, BluetoothChange.ADVERTISEMENT)
 
-async def test_bluetooth_confirm_lost_device(hass: HomeAssistant) -> None:
+    assert len(hass.config_entries.flow.async_progress_by_handler("improv_ble")) == 0
-    """Test bluetooth confirm step when device can no longer be connected to."""
-    result = await hass.config_entries.flow.async_init(
-        DOMAIN,
-        context={"source": config_entries.SOURCE_BLUETOOTH},
-        data=IMPROV_BLE_DISCOVERY_INFO,
-    )
-    assert result["type"] == FlowResultType.FORM
-    assert result["step_id"] == "bluetooth_confirm"
-    assert result["errors"] is None
-
-    with patch(
-        f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", return_value=False
-    ), patch(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
-        return_value=None,
-    ):
-        result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
-        assert result["type"] == FlowResultType.ABORT
-        assert result["reason"] == "cannot_connect"
 
 
 async def test_bluetooth_step_success(hass: HomeAssistant) -> None:
@@ -198,13 +159,9 @@ async def test_bluetooth_step_success(hass: HomeAssistant) -> None:
     assert result["step_id"] == "bluetooth_confirm"
     assert result["errors"] is None
 
-    with patch(
+    await _test_common_success_wo_identify(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
+        hass, result, IMPROV_BLE_DISCOVERY_INFO.address
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
+    )
-    ):
-        await _test_common_success_wo_identify(
-            hass, result, IMPROV_BLE_DISCOVERY_INFO.address
-        )
 
 
 async def test_bluetooth_step_success_identify(hass: HomeAssistant) -> None:
@@ -218,13 +175,9 @@ async def test_bluetooth_step_success_identify(hass: HomeAssistant) -> None:
     assert result["step_id"] == "bluetooth_confirm"
     assert result["errors"] is None
 
-    with patch(
+    await _test_common_success_with_identify(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
+        hass, result, IMPROV_BLE_DISCOVERY_INFO.address
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
+    )
-    ):
-        await _test_common_success_with_identify(
-            hass, result, IMPROV_BLE_DISCOVERY_INFO.address
-        )
 
 
 async def _test_common_success_with_identify(
@@ -442,9 +395,6 @@ async def test_can_identify_fails(hass: HomeAssistant, exc, error) -> None:
 
     with patch(
         f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", side_effect=exc
-    ), patch(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -480,9 +430,6 @@ async def test_identify_fails(hass: HomeAssistant, exc, error) -> None:
 
     with patch(
         f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", return_value=True
-    ), patch(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -526,9 +473,6 @@ async def test_need_authorization_fails(hass: HomeAssistant, exc, error) -> None
 
     with patch(
         f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", return_value=False
-    ), patch(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -573,9 +517,6 @@ async def test_authorize_fails(hass: HomeAssistant, exc, error) -> None:
 
     with patch(
         f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", return_value=False
-    ), patch(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -656,11 +597,7 @@ async def _test_provision_error(hass: HomeAssistant, exc) -> None:
 )
 async def test_provision_fails(hass: HomeAssistant, exc, error) -> None:
     """Test bluetooth flow with error."""
-    with patch(
+    flow_id = await _test_provision_error(hass, exc)
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
-    ):
-        flow_id = await _test_provision_error(hass, exc)
 
     result = await hass.config_entries.flow.async_configure(flow_id)
     assert result["type"] == FlowResultType.ABORT
@@ -683,9 +620,6 @@ async def test_provision_not_authorized(hass: HomeAssistant, exc, error) -> None
     with patch(
         f"{IMPROV_BLE}.config_flow.ImprovBLEClient.subscribe_state_updates",
         side_effect=subscribe_state_updates,
-    ), patch(
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
     ):
         flow_id = await _test_provision_error(hass, exc)
     result = await hass.config_entries.flow.async_configure(flow_id)
@@ -705,11 +639,7 @@ async def test_provision_not_authorized(hass: HomeAssistant, exc, error) -> None
 )
 async def test_provision_retry(hass: HomeAssistant, exc, error) -> None:
     """Test bluetooth flow with error."""
-    with patch(
+    flow_id = await _test_provision_error(hass, exc)
-        f"{IMPROV_BLE}.config_flow.async_last_service_info",
-        return_value=IMPROV_BLE_DISCOVERY_INFO,
-    ):
-        flow_id = await _test_provision_error(hass, exc)
 
     result = await hass.config_entries.flow.async_configure(flow_id)
     assert result["type"] == FlowResultType.FORM