Add flow chaining from Improv BLE to integration config flows (#154415)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

tests/components/esphome/test_config_flow.py

@@ -13,6 +13,7 @@ from aioesphomeapi import (
     InvalidEncryptionKeyAPIError,
     RequiresEncryptionAPIError,
     ResolveAPIError,
+    wifi_mac_to_bluetooth_mac,
 )
 import aiohttp
 import pytest
@@ -2817,3 +2818,79 @@ async def test_user_flow_zwave_discovery_aborts(
 
     # Verify next_flow was NOT set since Z-Wave flow aborted
     assert "next_flow" not in result
+
+
+@pytest.mark.usefixtures("mock_setup_entry", "mock_zeroconf")
+async def test_zeroconf_notifies_improv_ble(
+    hass: HomeAssistant,
+    mock_client: APIClient,
+) -> None:
+    """Test that zeroconf discovery notifies improv_ble integration."""
+    service_info = ZeroconfServiceInfo(
+        ip_address=ip_address("192.168.43.183"),
+        ip_addresses=[ip_address("192.168.43.183")],
+        hostname="test8266.local.",
+        name="mock_name",
+        port=6053,
+        properties={
+            "mac": "aabbccddeeff",
+        },
+        type="mock_type",
+    )
+
+    # Patch improv_ble to ensure it's available and track calls
+    with patch(
+        "homeassistant.components.improv_ble.async_register_next_flow"
+    ) as mock_register:
+        flow = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_ZEROCONF},
+            data=service_info,
+        )
+
+    assert flow["type"] is FlowResultType.FORM
+    assert flow["step_id"] == "discovery_confirm"
+
+    # Verify improv_ble.async_register_next_flow was called with correct parameters
+    assert len(mock_register.mock_calls) == 1
+    call_args = mock_register.mock_calls[0].args
+    assert call_args[0] is hass  # HomeAssistant instance
+    # WiFi MAC aabbccddeeff + 1 = Bluetooth MAC aabbccddee00
+    # (wifi_mac_to_bluetooth_mac from aioesphomeapi)
+    expected_ble_mac = wifi_mac_to_bluetooth_mac("aa:bb:cc:dd:ee:ff")
+    assert call_args[1] == expected_ble_mac  # BLE MAC address
+    assert call_args[2] == flow["flow_id"]  # Flow ID
+
+
+@pytest.mark.usefixtures("mock_setup_entry", "mock_zeroconf")
+async def test_zeroconf_when_improv_ble_not_available(
+    hass: HomeAssistant,
+    mock_client: APIClient,
+) -> None:
+    """Test that zeroconf discovery works when improv_ble is not available."""
+    service_info = ZeroconfServiceInfo(
+        ip_address=ip_address("192.168.43.183"),
+        ip_addresses=[ip_address("192.168.43.183")],
+        hostname="test8266.local.",
+        name="mock_name",
+        port=6053,
+        properties={
+            "mac": "aabbccddeeff",
+        },
+        type="mock_type",
+    )
+
+    # Mock async_import_module to return None (simulating improv_ble not available)
+    with patch(
+        "homeassistant.components.esphome.config_flow.async_import_module",
+        return_value=None,
+    ):
+        flow = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_ZEROCONF},
+            data=service_info,
+        )
+
+    # Flow should still work even without improv_ble
+    assert flow["type"] is FlowResultType.FORM
+    assert flow["step_id"] == "discovery_confirm"

tests/components/improv_ble/test_config_flow.py

@@ -1,19 +1,27 @@
 """Test the Improv via BLE config flow."""
 
+import asyncio
 from collections.abc import Callable
 from unittest.mock import patch
 
 from bleak.exc import BleakError
-from improv_ble_client import Error, State, errors as improv_ble_errors
+from improv_ble_client import (
+    SERVICE_DATA_UUID,
+    SERVICE_UUID,
+    Error,
+    State,
+    errors as improv_ble_errors,
+)
 import pytest
 
 from homeassistant import config_entries
+from homeassistant.components import improv_ble
 from homeassistant.components.bluetooth import (
     BluetoothChange,
     BluetoothServiceInfoBleak,
 )
 from homeassistant.components.improv_ble.const import DOMAIN
-from homeassistant.config_entries import SOURCE_IGNORE
+from homeassistant.config_entries import SOURCE_IGNORE, FlowType
 from homeassistant.const import CONF_ADDRESS
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResult, FlowResultType
@@ -34,6 +42,43 @@ from tests.components.bluetooth import (
 
 IMPROV_BLE = "homeassistant.components.improv_ble"
 
+# Discovery info for target flow devices (used for flow chaining tests)
+IMPROV_BLE_DISCOVERY_INFO_TARGET1 = BluetoothServiceInfoBleak(
+    name="target_device",
+    address="AA:BB:CC:DD:EE:F1",
+    rssi=-60,
+    manufacturer_data={},
+    service_uuids=[SERVICE_UUID],
+    service_data={SERVICE_DATA_UUID: b"\x01\x00\x00\x00\x00\x00"},
+    source="local",
+    device=generate_ble_device(address="AA:BB:CC:DD:EE:F1", name="target_device"),
+    advertisement=generate_advertisement_data(
+        service_uuids=[SERVICE_UUID],
+        service_data={SERVICE_DATA_UUID: b"\x01\x00\x00\x00\x00\x00"},
+    ),
+    time=0,
+    connectable=True,
+    tx_power=-127,
+)
+
+IMPROV_BLE_DISCOVERY_INFO_TARGET2 = BluetoothServiceInfoBleak(
+    name="esphome_device",
+    address="AA:BB:CC:DD:EE:F2",
+    rssi=-60,
+    manufacturer_data={},
+    service_uuids=[SERVICE_UUID],
+    service_data={SERVICE_DATA_UUID: b"\x01\x00\x00\x00\x00\x00"},
+    source="local",
+    device=generate_ble_device(address="AA:BB:CC:DD:EE:F2", name="esphome_device"),
+    advertisement=generate_advertisement_data(
+        service_uuids=[SERVICE_UUID],
+        service_data={SERVICE_DATA_UUID: b"\x01\x00\x00\x00\x00\x00"},
+    ),
+    time=0,
+    connectable=True,
+    tx_power=-127,
+)
+
 
 @pytest.mark.parametrize(
     ("url", "abort_reason", "placeholders"),
@@ -269,6 +314,7 @@ async def test_bluetooth_rediscovery_after_successful_provision(
             f"{IMPROV_BLE}.config_flow.ImprovBLEClient.provision",
             return_value=None,
         ),
+        patch(f"{IMPROV_BLE}.config_flow.PROVISIONING_TIMEOUT", 0.0000001),
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"], {"ssid": "TestNetwork", "password": "secret"}
@@ -394,6 +440,7 @@ async def _test_common_success(
     url: str | None = None,
     abort_reason: str = "provision_successful",
     placeholders: dict[str, str] | None = None,
+    patch_timeout_for_tests=None,
 ) -> None:
     """Test bluetooth and user flow success paths."""
 
@@ -406,6 +453,7 @@ async def _test_common_success(
             f"{IMPROV_BLE}.config_flow.ImprovBLEClient.provision",
             return_value=url,
         ) as mock_provision,
+        patch(f"{IMPROV_BLE}.config_flow.PROVISIONING_TIMEOUT", 0.0000001),
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"], {"ssid": "MyWIFI", "password": "secret"}
@@ -480,6 +528,7 @@ async def _test_common_success_w_authorize(
             f"{IMPROV_BLE}.config_flow.ImprovBLEClient.provision",
             return_value="http://blabla.local",
         ) as mock_provision,
+        patch(f"{IMPROV_BLE}.config_flow.PROVISIONING_TIMEOUT", 0.0000001),
     ):
         result = await hass.config_entries.flow.async_configure(result["flow_id"])
         assert result["type"] is FlowResultType.SHOW_PROGRESS
@@ -722,6 +771,7 @@ async def _test_provision_error(hass: HomeAssistant, exc) -> str:
             f"{IMPROV_BLE}.config_flow.ImprovBLEClient.provision",
             side_effect=exc,
         ),
+        patch(f"{IMPROV_BLE}.config_flow.PROVISIONING_TIMEOUT", 0.0000001),
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"], {"ssid": "MyWIFI", "password": "secret"}
@@ -812,6 +862,275 @@ async def test_provision_fails_invalid_data(
     )
 
 
+async def test_flow_chaining_with_next_flow(hass: HomeAssistant) -> None:
+    """Test flow chaining when another integration registers a next flow."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+        data=IMPROV_BLE_DISCOVERY_INFO,
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "bluetooth_confirm"
+
+    # Confirm bluetooth setup
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "bluetooth_confirm"
+
+    # Start provisioning
+    with patch(
+        f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", return_value=False
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_ADDRESS: IMPROV_BLE_DISCOVERY_INFO.address},
+        )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "provision"
+
+    with (
+        patch(
+            f"{IMPROV_BLE}.config_flow.ImprovBLEClient.need_authorization",
+            return_value=False,
+        ),
+        patch(
+            f"{IMPROV_BLE}.config_flow.ImprovBLEClient.provision",
+            return_value=None,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {"ssid": "TestNetwork", "password": "secret"}
+        )
+        assert result["type"] is FlowResultType.SHOW_PROGRESS
+        assert result["progress_action"] == "provisioning"
+        assert result["step_id"] == "do_provision"
+
+        # Yield to allow the background task to create the future
+        await asyncio.sleep(0)  # task is created with eager_start=False
+
+        # Create a dummy target flow using a different device address
+        target_result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_BLUETOOTH},
+            data=IMPROV_BLE_DISCOVERY_INFO_TARGET1,
+        )
+        next_config_flow_id = target_result["flow_id"]
+
+        # Simulate another integration discovering the device and registering a flow
+        # This happens while provision is waiting on the future
+        improv_ble.async_register_next_flow(
+            hass, IMPROV_BLE_DISCOVERY_INFO.address, next_config_flow_id
+        )
+
+        await hass.async_block_till_done()
+
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "provision_successful"
+    assert result["next_flow"] == (FlowType.CONFIG_FLOW, next_config_flow_id)
+
+
+async def test_flow_chaining_timeout(hass: HomeAssistant) -> None:
+    """Test flow chaining timeout when no integration discovers the device."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+        data=IMPROV_BLE_DISCOVERY_INFO,
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "bluetooth_confirm"
+
+    # Confirm bluetooth setup
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "bluetooth_confirm"
+
+    # Start provisioning
+    with patch(
+        f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", return_value=False
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_ADDRESS: IMPROV_BLE_DISCOVERY_INFO.address},
+        )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "provision"
+
+    # Complete provisioning successfully but no integration discovers the device
+    with (
+        patch(
+            f"{IMPROV_BLE}.config_flow.ImprovBLEClient.need_authorization",
+            return_value=False,
+        ),
+        patch(
+            f"{IMPROV_BLE}.config_flow.ImprovBLEClient.provision",
+            return_value=None,
+        ),
+        patch("asyncio.wait_for", side_effect=TimeoutError),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {"ssid": "TestNetwork", "password": "secret"}
+        )
+        assert result["type"] is FlowResultType.SHOW_PROGRESS
+        assert result["progress_action"] == "provisioning"
+        assert result["step_id"] == "do_provision"
+        await hass.async_block_till_done()
+
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "provision_successful"
+    assert "next_flow" not in result
+
+
+async def test_flow_chaining_with_redirect_url(hass: HomeAssistant) -> None:
+    """Test flow chaining takes precedence over redirect URL."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+        data=IMPROV_BLE_DISCOVERY_INFO,
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "bluetooth_confirm"
+
+    # Confirm bluetooth setup
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "bluetooth_confirm"
+
+    # Start provisioning
+    with patch(
+        f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", return_value=False
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_ADDRESS: IMPROV_BLE_DISCOVERY_INFO.address},
+        )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "provision"
+
+    with (
+        patch(
+            f"{IMPROV_BLE}.config_flow.ImprovBLEClient.need_authorization",
+            return_value=False,
+        ),
+        patch(
+            f"{IMPROV_BLE}.config_flow.ImprovBLEClient.provision",
+            return_value="http://blabla.local",
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {"ssid": "TestNetwork", "password": "secret"}
+        )
+        assert result["type"] is FlowResultType.SHOW_PROGRESS
+        assert result["progress_action"] == "provisioning"
+        assert result["step_id"] == "do_provision"
+
+        # Yield to allow the background task to create the future
+        await asyncio.sleep(0)  # task is created with eager_start=False
+
+        # Create a dummy target flow using a different device address
+        target_result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_BLUETOOTH},
+            data=IMPROV_BLE_DISCOVERY_INFO_TARGET2,
+        )
+        esphome_flow_id = target_result["flow_id"]
+
+        # Simulate ESPHome discovering the device and notifying Improv BLE
+        # This happens while provision is still running
+        improv_ble.async_register_next_flow(
+            hass, IMPROV_BLE_DISCOVERY_INFO.address, esphome_flow_id
+        )
+
+        await hass.async_block_till_done()
+
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    assert result["type"] is FlowResultType.ABORT
+    # Should use next_flow instead of redirect URL
+    assert result["reason"] == "provision_successful"
+    assert result["next_flow"] == (FlowType.CONFIG_FLOW, esphome_flow_id)
+
+
+async def test_flow_chaining_future_already_done(
+    hass: HomeAssistant, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test async_register_next_flow when future is already done."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+        data=IMPROV_BLE_DISCOVERY_INFO,
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "bluetooth_confirm"
+
+    # Confirm bluetooth setup
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "bluetooth_confirm"
+
+    # Start provisioning
+    with patch(
+        f"{IMPROV_BLE}.config_flow.ImprovBLEClient.can_identify", return_value=False
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_ADDRESS: IMPROV_BLE_DISCOVERY_INFO.address},
+        )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "provision"
+
+    with (
+        patch(
+            f"{IMPROV_BLE}.config_flow.ImprovBLEClient.need_authorization",
+            return_value=False,
+        ),
+        patch(
+            f"{IMPROV_BLE}.config_flow.ImprovBLEClient.provision",
+            return_value=None,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {"ssid": "TestNetwork", "password": "secret"}
+        )
+        assert result["type"] is FlowResultType.SHOW_PROGRESS
+        assert result["progress_action"] == "provisioning"
+        assert result["step_id"] == "do_provision"
+
+        # Yield to allow the background task to create the future
+        await asyncio.sleep(0)  # task is created with eager_start=False
+
+        # Create a target flow for the first call
+        target_result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_BLUETOOTH},
+            data=IMPROV_BLE_DISCOVERY_INFO_TARGET1,
+        )
+        first_flow_id = target_result["flow_id"]
+
+        # First call resolves the future
+        improv_ble.async_register_next_flow(
+            hass, IMPROV_BLE_DISCOVERY_INFO.address, first_flow_id
+        )
+
+        # Second call immediately after - future is now done but still in registry
+        # This call should be ignored with a debug log
+        caplog.clear()
+        improv_ble.async_register_next_flow(
+            hass, IMPROV_BLE_DISCOVERY_INFO.address, "second_flow_id"
+        )
+
+        # Verify the debug log message was emitted
+        assert "Future for aa:bb:cc:dd:ee:f0 already done" in caplog.text
+        assert "ignoring flow_id second_flow_id" in caplog.text
+
+        await hass.async_block_till_done()
+
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "provision_successful"
+    assert result["next_flow"] == (FlowType.CONFIG_FLOW, first_flow_id)
+
+
 async def test_bluetooth_name_update(hass: HomeAssistant) -> None:
     """Test that discovery notification title updates when device name changes."""
     with patch(