From ca31134caa32fe41467d40fd0ae33b383d41c7a6 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 7 Dec 2025 14:28:53 -0600 Subject: [PATCH] Keep persistent BLE connection during Shelly WiFi provisioning (#158145) Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../components/shelly/config_flow.py | 155 ++- tests/components/shelly/test_config_flow.py | 1040 ++++++++++------- 2 files changed, 758 insertions(+), 437 deletions(-) diff --git a/homeassistant/components/shelly/config_flow.py b/homeassistant/components/shelly/config_flow.py index 7575b30f951..7c3cb74bf54 100644 --- a/homeassistant/components/shelly/config_flow.py +++ b/homeassistant/components/shelly/config_flow.py @@ -13,11 +13,6 @@ from aioshelly.ble.manufacturer_data import ( has_rpc_over_ble, parse_shelly_manufacturer_data, ) -from aioshelly.ble.provisioning import ( - async_provision_wifi, - async_scan_wifi_networks, - ble_rpc_device, -) from aioshelly.block_device import BlockDevice from aioshelly.common import ConnectionOptions, get_info from aioshelly.const import BLOCK_GENERATIONS, DEFAULT_HTTP_PORT, RPC_GENERATIONS @@ -119,31 +114,6 @@ MANUAL_ENTRY_STRING = "manual" DISCOVERY_SOURCES = {SOURCE_BLUETOOTH, SOURCE_ZEROCONF} -async def async_get_ip_from_ble(ble_device: BLEDevice) -> str | None: - """Get device IP address via BLE after WiFi provisioning. - - Args: - ble_device: BLE device to query - - Returns: - IP address string if available, None otherwise - - """ - try: - async with ble_rpc_device(ble_device) as device: - await device.update_status() - if ( - (wifi := device.status.get("wifi")) - and isinstance(wifi, dict) - and (ip := wifi.get("sta_ip")) - ): - return cast(str, ip) - return None - except (DeviceConnectionError, RpcCallError) as err: - LOGGER.debug("Failed to get IP via BLE: %s", err) - return None - - # BLE provisioning flow steps that are in the finishing state # Used to determine if a BLE flow should be aborted when zeroconf discovers the device BLUETOOTH_FINISHING_STEPS = {"do_provision", "provision_done"} @@ -252,6 +222,7 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): disable_ap_after_provision: bool = True disable_ble_rpc_after_provision: bool = True _discovered_devices: dict[str, DiscoveredDeviceZeroconf | DiscoveredDeviceBluetooth] + _ble_rpc_device: RpcDevice | None = None @staticmethod def _get_name_from_mac_and_ble_model( @@ -300,6 +271,81 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): return mac, device_name + async def _async_ensure_ble_connected(self) -> RpcDevice: + """Ensure BLE RPC device is connected, reconnecting if needed. + + Maintains a persistent BLE connection across config flow steps to avoid + the overhead of reconnecting between WiFi scan and provisioning steps. + + Returns: + Connected RpcDevice instance + + Raises: + DeviceConnectionError: If connection fails + RpcCallError: If ping fails after connection + + """ + if TYPE_CHECKING: + assert self.ble_device is not None + + if self._ble_rpc_device is not None and self._ble_rpc_device.connected: + # Ping to verify connection is still alive + try: + await self._ble_rpc_device.update_status() + except (DeviceConnectionError, RpcCallError): + # Connection dropped, need to reconnect + LOGGER.debug("BLE connection lost, reconnecting") + await self._async_disconnect_ble() + else: + return self._ble_rpc_device + + # Create new connection + LOGGER.debug("Creating new BLE RPC connection to %s", self.ble_device.address) + options = ConnectionOptions(ble_device=self.ble_device) + device = await RpcDevice.create( + aiohttp_session=None, ws_context=None, ip_or_options=options + ) + try: + await device.initialize() + except (DeviceConnectionError, RpcCallError): + await device.shutdown() + raise + self._ble_rpc_device = device + return self._ble_rpc_device + + async def _async_disconnect_ble(self) -> None: + """Disconnect and cleanup BLE RPC device.""" + if self._ble_rpc_device is not None: + try: + await self._ble_rpc_device.shutdown() + except Exception: # noqa: BLE001 + LOGGER.debug("Error during BLE shutdown", exc_info=True) + finally: + self._ble_rpc_device = None + + async def _async_get_ip_from_ble(self) -> str | None: + """Get device IP address via BLE after WiFi provisioning. + + Uses the persistent BLE connection to get the device's sta_ip from status. + + Returns: + IP address string if available, None otherwise + + """ + try: + device = await self._async_ensure_ble_connected() + except (DeviceConnectionError, RpcCallError) as err: + LOGGER.debug("Failed to get IP via BLE: %s", err) + return None + + if ( + (wifi := device.status.get("wifi")) + and isinstance(wifi, dict) + and (ip := wifi.get("sta_ip")) + ): + return cast(str, ip) + return None + async def _async_discover_zeroconf_devices( self, ) -> dict[str, DiscoveredDeviceZeroconf]: @@ -737,20 +783,21 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): password = user_input[CONF_PASSWORD] return await self.async_step_do_provision({"password": password}) - # Scan for WiFi networks via BLE - if TYPE_CHECKING: - assert self.ble_device is not None + # Scan for WiFi networks via BLE using persistent connection try: - self.wifi_networks = await async_scan_wifi_networks(self.ble_device) + device = await self._async_ensure_ble_connected() + self.wifi_networks = await device.wifi_scan() except (DeviceConnectionError, RpcCallError) as err: LOGGER.debug("Failed to scan WiFi networks via BLE: %s", err) # "Writing is not permitted" error means device rejects BLE writes # and BLE provisioning is disabled - user must use Shelly app if "not permitted" in str(err): + await self._async_disconnect_ble() return self.async_abort(reason="ble_not_permitted") return await self.async_step_wifi_scan_failed() except Exception: # noqa: BLE001 LOGGER.exception("Unexpected exception during WiFi scan") + await self._async_disconnect_ble() return self.async_abort(reason="unknown") # Sort by RSSI (strongest signal first - higher/less negative values first) @@ -871,17 +918,21 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): Returns the flow result to be stored in self._provision_result, or None if failed. """ - # Provision WiFi via BLE - if TYPE_CHECKING: - assert self.ble_device is not None + # Provision WiFi via BLE using persistent connection try: - await async_provision_wifi(self.ble_device, self.selected_ssid, password) + device = await self._async_ensure_ble_connected() + await device.wifi_setconfig( + sta_ssid=self.selected_ssid, + sta_password=password, + sta_enable=True, + ) except (DeviceConnectionError, RpcCallError) as err: LOGGER.debug("Failed to provision WiFi via BLE: %s", err) # BLE connection/communication failed - allow retry from network selection return None except Exception: # noqa: BLE001 LOGGER.exception("Unexpected exception during WiFi provisioning") + await self._async_disconnect_ble() return self.async_abort(reason="unknown") LOGGER.debug( @@ -919,7 +970,7 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): LOGGER.debug( "Active lookup failed, trying to get IP address via BLE as fallback" ) - if ip := await async_get_ip_from_ble(self.ble_device): + if ip := await self._async_get_ip_from_ble(): LOGGER.debug("Got IP %s from BLE, using it", ip) state.host = ip state.port = DEFAULT_HTTP_PORT @@ -996,12 +1047,17 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): if TYPE_CHECKING: assert mac is not None - async with self._async_provision_context(mac) as state: - self._provision_result = ( - await self._async_provision_wifi_and_wait_for_zeroconf( - mac, password, state + try: + async with self._async_provision_context(mac) as state: + self._provision_result = ( + await self._async_provision_wifi_and_wait_for_zeroconf( + mac, password, state + ) ) - ) + finally: + # Always disconnect BLE after provisioning attempt completes + # We either succeeded (and will use IP now) or failed (and user will retry) + await self._async_disconnect_ble() async def async_step_do_provision( self, user_input: dict[str, Any] | None = None @@ -1220,6 +1276,17 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): """Get info from shelly device.""" return await get_info(async_get_clientsession(self.hass), host, port=port) + @callback + def async_remove(self) -> None: + """Handle flow removal - cleanup BLE connection.""" + super().async_remove() + if self._ble_rpc_device is not None: + # Schedule cleanup as background task since async_remove is sync + self.hass.async_create_background_task( + self._async_disconnect_ble(), + name="shelly_config_flow_ble_cleanup", + ) + @staticmethod @callback def async_get_options_flow(config_entry: ShellyConfigEntry) -> OptionsFlowHandler: diff --git a/tests/components/shelly/test_config_flow.py b/tests/components/shelly/test_config_flow.py index 5741e95f16c..2fd0eb71d5d 100644 --- a/tests/components/shelly/test_config_flow.py +++ b/tests/components/shelly/test_config_flow.py @@ -312,24 +312,55 @@ def mock_zeroconf_async_get_instance() -> Generator[AsyncMock]: yield mock_aiozc -@pytest.fixture -def mock_wifi_scan() -> Generator[AsyncMock]: - """Mock async_scan_wifi_networks.""" - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - new=AsyncMock(return_value=[{"ssid": "TestNetwork", "rssi": -50, "auth": 2}]), - ) as mock_scan: - yield mock_scan +def create_mock_ble_rpc_device( + wifi_networks: list[dict[str, Any]] | None = None, + sta_ip: str = "192.168.1.100", +) -> AsyncMock: + """Create a mock BLE RPC device for provisioning tests.""" + if wifi_networks is None: + wifi_networks = [{"ssid": "TestNetwork", "rssi": -50, "auth": 2}] + mock_device = AsyncMock() + mock_device.initialize = AsyncMock() + mock_device.shutdown = AsyncMock() + mock_device.connected = True + mock_device.call_rpc = AsyncMock(return_value={}) + mock_device.wifi_scan = AsyncMock(return_value=wifi_networks) + mock_device.wifi_setconfig = AsyncMock(return_value={}) + mock_device.update_status = AsyncMock() + mock_device.status = {"wifi": {"sta_ip": sta_ip}} + return mock_device @pytest.fixture -def mock_wifi_provision() -> Generator[AsyncMock]: - """Mock async_provision_wifi.""" - with patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - new=AsyncMock(), - ) as mock_provision: - yield mock_provision +def mock_ble_rpc_device() -> AsyncMock: + """Create a mock BLE RPC device for provisioning tests.""" + return create_mock_ble_rpc_device() + + +@pytest.fixture +def mock_ble_rpc_device_class(mock_ble_rpc_device: AsyncMock) -> Generator[MagicMock]: + """Mock RpcDevice.create for BLE provisioning. + + This fixture patches RpcDevice to return a mock for BLE connections + (where aiohttp_session is None) while allowing normal behavior or + other patches for IP-based connections. + """ + original_rpc_device = config_flow.RpcDevice + + async def mock_create( + aiohttp_session: Any, ws_context: Any, ip_or_options: Any + ) -> Any: + # BLE connections have aiohttp_session=None + if aiohttp_session is None: + return mock_ble_rpc_device + # For IP connections, use the original class (or let other patches handle it) + return await original_rpc_device.create( + aiohttp_session, ws_context, ip_or_options + ) + + with patch("homeassistant.components.shelly.config_flow.RpcDevice") as mock_class: + mock_class.create = AsyncMock(side_effect=mock_create) + yield mock_class @pytest.fixture(autouse=True) @@ -1092,11 +1123,18 @@ async def test_user_flow_both_ble_and_zeroconf_prefers_zeroconf( assert result["data"][CONF_PORT] == 80 +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_user_flow_with_ble_devices( hass: HomeAssistant, mock_discovery: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test user flow shows discovered BLE devices.""" + # Configure mock BLE device for this test + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "TestNetwork", "rssi": -50, "auth": 2} + ] + # Mock empty zeroconf discovery mock_discovery.return_value = [] @@ -1132,7 +1170,7 @@ async def test_user_flow_with_ble_devices( flows = hass.config_entries.flow.async_progress_by_handler(DOMAIN) for flow in flows: if flow["context"]["source"] == config_entries.SOURCE_BLUETOOTH: - await hass.config_entries.flow.async_abort(flow["flow_id"]) + hass.config_entries.flow.async_abort(flow["flow_id"]) result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_USER} @@ -1164,19 +1202,14 @@ async def test_user_flow_with_ble_devices( assert result["type"] is FlowResultType.FORM assert result["step_id"] == "bluetooth_confirm" - # Complete WiFi provisioning - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "TestNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {}, - ) + # Confirm BLE provisioning - wifi_scan is handled by fixture + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) # Select network and enter WiFi credentials to complete with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("192.168.1.100", 80), @@ -1190,10 +1223,6 @@ async def test_user_flow_with_ble_devices( "gen": 2, }, ), - patch( - "homeassistant.components.shelly.config_flow.RpcDevice.create", - return_value=create_mock_rpc_device("Test Device"), - ), ): result = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -1799,11 +1828,18 @@ async def test_user_flow_select_zeroconf_device_not_fully_provisioned( assert result["reason"] == "firmware_not_fully_provisioned" +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_user_flow_select_ble_device( hass: HomeAssistant, mock_discovery: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test selecting a BLE device goes to provisioning flow.""" + # Configure mock BLE device for this test + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Mock empty zeroconf discovery mock_discovery.return_value = [] @@ -1830,24 +1866,17 @@ async def test_user_flow_select_ble_device( assert result["type"] is FlowResultType.FORM assert result["step_id"] == "bluetooth_confirm" - # Confirm BLE provisioning and scan for WiFi networks - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[ - {"ssid": "MyNetwork", "rssi": -50, "auth": 2}, - ], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {}, - ) + # Confirm BLE provisioning - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) assert result["type"] is FlowResultType.FORM assert result["step_id"] == "wifi_scan" # Select network and enter password to provision with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("192.168.1.200", 80), @@ -1861,10 +1890,6 @@ async def test_user_flow_select_ble_device( "gen": 2, }, ), - patch( - "homeassistant.components.shelly.config_flow.RpcDevice.create", - return_value=create_mock_rpc_device("Test BLE Device"), - ), ): result = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -1883,7 +1908,7 @@ async def test_user_flow_select_ble_device( # Should create entry assert result["type"] is FlowResultType.CREATE_ENTRY assert result["result"].unique_id == "CCBA97C2D670" - assert result["title"] == "Test BLE Device" + assert result["title"] == "Test name" async def test_user_flow_filters_devices_with_active_discovery_flows( @@ -3185,13 +3210,21 @@ async def test_zeroconf_wrong_device_name( assert len(mock_setup_entry.mock_calls) == 1 -@pytest.mark.usefixtures("mock_rpc_device", "mock_zeroconf") +@pytest.mark.usefixtures( + "mock_rpc_device", "mock_zeroconf", "mock_ble_rpc_device_class" +) async def test_bluetooth_discovery( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test bluetooth discovery and complete provisioning.""" + # Configure mock BLE device for this test + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -3205,18 +3238,11 @@ async def test_bluetooth_discovery( assert result["step_id"] == "bluetooth_confirm" assert result["description_placeholders"]["name"] == "ShellyPlus2PM-C049EF8873E8" - # Confirm - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # Confirm - BLE device will be used for wifi_scan via fixture + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) # Select network and enter password to provision with ( - patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - ), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -3253,12 +3279,19 @@ async def test_bluetooth_discovery( assert len(mock_setup_entry.mock_calls) == 1 +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provisioning_clears_match_history( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test bluetooth provisioning clears match history at discovery start and after successful provisioning.""" + # Configure mock BLE device for this test + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO_FOR_CLEAR_TEST) @@ -3274,23 +3307,14 @@ async def test_bluetooth_provisioning_clears_match_history( assert result["type"] is FlowResultType.FORM assert result["step_id"] == "bluetooth_confirm" - # Confirm - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], {} - ) + # Confirm - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) # Reset mock to only count calls during provisioning mock_clear.reset_mock() # Select network and enter password to provision with ( - patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - ), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -3343,12 +3367,19 @@ async def test_bluetooth_discovery_no_rpc_over_ble( assert result["reason"] == "invalid_discovery_info" +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_factory_reset_rediscovery( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test device can be rediscovered after factory reset when RPC-over-BLE is re-enabled.""" + # Configure mock BLE device for this test + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # First discovery: device is already provisioned (no RPC-over-BLE) # Inject the device without RPC so it's in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO_NO_RPC) @@ -3382,17 +3413,11 @@ async def test_bluetooth_factory_reset_rediscovery( result["context"]["title_placeholders"]["name"] == "ShellyPlus2PM-C049EF8873E8" ) - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # Confirm - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) # Select network and enter password to provision with ( - patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - ), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -3561,13 +3586,22 @@ async def test_bluetooth_discovery_no_ble_device( assert result["reason"] == "cannot_connect" -@pytest.mark.usefixtures("mock_rpc_device", "mock_zeroconf") +@pytest.mark.usefixtures( + "mock_rpc_device", "mock_zeroconf", "mock_ble_rpc_device_class" +) async def test_bluetooth_wifi_scan_success( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test WiFi scan via BLE.""" + # Configure mock BLE device for this test + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "Network1", "rssi": -50, "auth": 2}, + {"ssid": "Network2", "rssi": -60, "auth": 3}, + ] + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -3577,18 +3611,11 @@ async def test_bluetooth_wifi_scan_success( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm BLE provisioning and trigger wifi scan - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[ - {"ssid": "Network1", "rssi": -50, "auth": 2}, - {"ssid": "Network2", "rssi": -60, "auth": 3}, - ], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {}, - ) + # Confirm BLE provisioning - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) assert result["type"] is FlowResultType.FORM assert result["step_id"] == "wifi_scan" @@ -3597,9 +3624,6 @@ async def test_bluetooth_wifi_scan_success( # Select network and enter password to complete flow with ( - patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - ), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -3635,13 +3659,19 @@ async def test_bluetooth_wifi_scan_success( assert len(mock_setup_entry.mock_calls) == 1 -@pytest.mark.usefixtures("mock_rpc_device", "mock_zeroconf") +@pytest.mark.usefixtures( + "mock_rpc_device", "mock_zeroconf", "mock_ble_rpc_device_class" +) async def test_bluetooth_wifi_scan_failure( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test WiFi scan failure via BLE.""" + # Configure mock BLE device to fail first, then succeed + mock_ble_rpc_device.wifi_scan.side_effect = DeviceConnectionError + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -3651,35 +3681,32 @@ async def test_bluetooth_wifi_scan_failure( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm and trigger wifi scan that fails - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - side_effect=DeviceConnectionError, - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {}, - ) + # Confirm - wifi scan will fail + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) assert result["type"] is FlowResultType.FORM assert result["step_id"] == "wifi_scan_failed" - # Test retry and complete flow - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "Network1", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {}, - ) + # Now configure success for retry + mock_ble_rpc_device.wifi_scan.side_effect = None + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "Network1", "rssi": -50, "auth": 2} + ] + + # Test retry + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) assert result["type"] is FlowResultType.FORM assert result["step_id"] == "wifi_scan" # Select network and enter password to complete provisioning with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -3688,10 +3715,6 @@ async def test_bluetooth_wifi_scan_failure( "homeassistant.components.shelly.config_flow.get_info", return_value=MOCK_DEVICE_INFO, ), - patch( - "homeassistant.components.shelly.config_flow.RpcDevice.create", - return_value=create_mock_rpc_device("Test name"), - ), ): result = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -3719,11 +3742,19 @@ async def test_bluetooth_wifi_scan_failure( assert len(mock_setup_entry.mock_calls) == 1 -@pytest.mark.usefixtures("mock_rpc_device", "mock_zeroconf") +@pytest.mark.usefixtures( + "mock_rpc_device", "mock_zeroconf", "mock_ble_rpc_device_class" +) async def test_bluetooth_wifi_scan_ble_not_permitted( hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test WiFi scan when BLE is not permitted (cloud bound device).""" + # Configure mock BLE device to fail with permission error + mock_ble_rpc_device.wifi_scan.side_effect = DeviceConnectionError( + "Writing is not permitted" + ) + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -3733,27 +3764,31 @@ async def test_bluetooth_wifi_scan_ble_not_permitted( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm and trigger wifi scan that fails with permission error - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - side_effect=DeviceConnectionError("Writing is not permitted"), - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {}, - ) + # Confirm - wifi scan will fail with permission error + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) assert result["type"] is FlowResultType.ABORT assert result["reason"] == "ble_not_permitted" -@pytest.mark.usefixtures("mock_rpc_device", "mock_zeroconf") +@pytest.mark.usefixtures( + "mock_rpc_device", "mock_zeroconf", "mock_ble_rpc_device_class" +) async def test_bluetooth_wifi_credentials_and_provision_success( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test successful WiFi provisioning via BLE.""" + # Configure mock BLE device for this test + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -3763,28 +3798,17 @@ async def test_bluetooth_wifi_credentials_and_provision_success( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm BLE provisioning and scan for WiFi networks - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[ - {"ssid": "MyNetwork", "rssi": -50, "auth": 2}, - ], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {}, - ) + # Confirm BLE provisioning - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) assert result["type"] is FlowResultType.FORM assert result["step_id"] == "wifi_scan" # Select network and enter password to provision - mock_device = create_mock_rpc_device("Test name") - with ( - patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - ) as mock_provision, patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -3793,10 +3817,6 @@ async def test_bluetooth_wifi_credentials_and_provision_success( "homeassistant.components.shelly.config_flow.get_info", return_value=MOCK_DEVICE_INFO, ), - patch( - "homeassistant.components.shelly.config_flow.RpcDevice.create", - return_value=mock_device, - ), ): result = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -3823,18 +3843,28 @@ async def test_bluetooth_wifi_credentials_and_provision_success( CONF_SLEEP_PERIOD: 0, CONF_GEN: 2, } - assert mock_provision.call_count == 1 + mock_ble_rpc_device.wifi_setconfig.assert_called_once() assert len(mock_setup.mock_calls) == 1 assert len(mock_setup_entry.mock_calls) == 1 -@pytest.mark.usefixtures("mock_rpc_device", "mock_zeroconf") +@pytest.mark.usefixtures( + "mock_rpc_device", "mock_zeroconf", "mock_ble_rpc_device_class" +) async def test_bluetooth_wifi_provision_failure( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test WiFi provisioning failure via BLE.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # First provisioning attempt fails + mock_ble_rpc_device.wifi_setconfig.side_effect = DeviceConnectionError + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -3844,23 +3874,13 @@ async def test_bluetooth_wifi_provision_failure( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm and scan - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # Confirm and scan - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) - # Provision fails - with ( - patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - side_effect=DeviceConnectionError, - ), - patch( - "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", - return_value=None, - ), + # Provision fails - wifi_setconfig will fail + with patch( + "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", + return_value=None, ): result = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -3877,21 +3897,18 @@ async def test_bluetooth_wifi_provision_failure( assert result["type"] is FlowResultType.FORM assert result["step_id"] == "provision_failed" - # Test retry - go back to wifi scan and complete successfully - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # Reset wifi_setconfig for retry - now it succeeds + mock_ble_rpc_device.wifi_setconfig.side_effect = None + mock_ble_rpc_device.wifi_setconfig.return_value = {} + + # Test retry - go back to wifi scan + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) assert result["type"] is FlowResultType.FORM assert result["step_id"] == "wifi_scan" # Provision succeeds this time with ( - patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - ), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -3927,10 +3944,15 @@ async def test_bluetooth_wifi_provision_failure( assert len(mock_setup_entry.mock_calls) == 1 +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_wifi_scan_unexpected_exception( hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test unexpected exception during WiFi scan.""" + # Configure mock BLE device to raise unexpected exception + mock_ble_rpc_device.wifi_scan.side_effect = RuntimeError("Unexpected error") + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -3940,25 +3962,28 @@ async def test_bluetooth_wifi_scan_unexpected_exception( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm and trigger wifi scan that raises unexpected exception - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - side_effect=RuntimeError("Unexpected error"), - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {}, - ) + # Confirm - wifi scan will raise unexpected exception + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) assert result["type"] is FlowResultType.ABORT assert result["reason"] == "unknown" -@pytest.mark.usefixtures("mock_zeroconf") +@pytest.mark.usefixtures("mock_zeroconf", "mock_ble_rpc_device_class") async def test_bluetooth_provision_unexpected_exception( hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test unexpected exception during provisioning.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + mock_ble_rpc_device.wifi_setconfig.side_effect = RuntimeError("Unexpected error") + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -3968,23 +3993,13 @@ async def test_bluetooth_provision_unexpected_exception( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm and scan - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # Confirm and scan - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) # Provision raises unexpected exception in background task - with ( - patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - side_effect=RuntimeError("Unexpected error"), - ), - patch( - "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", - return_value=None, - ), + with patch( + "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", + return_value=None, ): result = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -4002,11 +4017,19 @@ async def test_bluetooth_provision_unexpected_exception( assert result["reason"] == "unknown" -@pytest.mark.usefixtures("mock_rpc_device", "mock_zeroconf") +@pytest.mark.usefixtures( + "mock_rpc_device", "mock_zeroconf", "mock_ble_rpc_device_class" +) async def test_bluetooth_provision_device_connection_error_after_wifi( hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test device connection error after WiFi provisioning.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -4016,16 +4039,11 @@ async def test_bluetooth_provision_device_connection_error_after_wifi( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm and scan - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # Confirm and scan - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) # Provision but get_info fails with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -4051,23 +4069,28 @@ async def test_bluetooth_provision_device_connection_error_after_wifi( assert result["step_id"] == "provision_failed" # User retries but BLE device raises unhandled exception - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - side_effect=RuntimeError("BLE device unavailable"), - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + mock_ble_rpc_device.wifi_scan.side_effect = RuntimeError("BLE device unavailable") + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) assert result["type"] is FlowResultType.ABORT assert result["reason"] == "unknown" -@pytest.mark.usefixtures("mock_rpc_device", "mock_zeroconf") +@pytest.mark.usefixtures( + "mock_rpc_device", "mock_zeroconf", "mock_ble_rpc_device_class" +) async def test_bluetooth_provision_requires_auth( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test device requires authentication after WiFi provisioning.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -4077,16 +4100,11 @@ async def test_bluetooth_provision_requires_auth( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm and scan - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # Confirm and scan - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) # Provision but device requires auth with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -4139,11 +4157,18 @@ async def test_bluetooth_provision_requires_auth( assert len(mock_setup_entry.mock_calls) == 1 -@pytest.mark.usefixtures("mock_zeroconf") +@pytest.mark.usefixtures("mock_zeroconf", "mock_ble_rpc_device_class") async def test_bluetooth_provision_validate_input_fails( hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, + mock_ble_rpc_device_class: MagicMock, ) -> None: """Test validate_input fails after WiFi provisioning.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -4153,16 +4178,24 @@ async def test_bluetooth_provision_validate_input_fails( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm and scan - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # Confirm and scan - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) - # Provision but validate_input fails + # Configure the mock to fail for IP-based connections (non-BLE) + # by making it raise DeviceConnectionError for any subsequent calls + original_create = mock_ble_rpc_device_class.create.side_effect + + async def fail_on_ip_connection(aiohttp_session, ws_context, ip_or_options): + # BLE connections have aiohttp_session=None + if aiohttp_session is None: + return mock_ble_rpc_device + # IP connections should fail + raise DeviceConnectionError("Simulated IP connection failure") + + mock_ble_rpc_device_class.create.side_effect = fail_on_ip_connection + + # Provision but validate_input fails (RpcDevice.create for IP connection fails) with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -4171,10 +4204,6 @@ async def test_bluetooth_provision_validate_input_fails( "homeassistant.components.shelly.config_flow.get_info", return_value=MOCK_DEVICE_INFO, ), - patch( - "homeassistant.components.shelly.config_flow.RpcDevice.create", - side_effect=DeviceConnectionError, - ), ): result = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -4191,22 +4220,28 @@ async def test_bluetooth_provision_validate_input_fails( assert result["type"] is FlowResultType.FORM assert result["step_id"] == "provision_failed" + # Restore original side_effect for retry + mock_ble_rpc_device_class.create.side_effect = original_create + # User retries but BLE device raises unhandled exception - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - side_effect=RuntimeError("BLE device unavailable"), - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + mock_ble_rpc_device.wifi_scan.side_effect = RuntimeError("BLE device unavailable") + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) assert result["type"] is FlowResultType.ABORT assert result["reason"] == "unknown" -@pytest.mark.usefixtures("mock_zeroconf") +@pytest.mark.usefixtures("mock_zeroconf", "mock_ble_rpc_device_class") async def test_bluetooth_provision_firmware_not_fully_provisioned( hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test device firmware not fully provisioned after WiFi provisioning.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Inject BLE device so it's available in the bluetooth scanner inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -4216,16 +4251,11 @@ async def test_bluetooth_provision_firmware_not_fully_provisioned( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Confirm and scan - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # Confirm and scan - wifi_scan handled by fixture + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) # Provision but device has no model (firmware not fully provisioned) with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -4255,13 +4285,19 @@ async def test_bluetooth_provision_firmware_not_fully_provisioned( assert result["reason"] == "firmware_not_fully_provisioned" -@pytest.mark.usefixtures("mock_zeroconf") +@pytest.mark.usefixtures("mock_zeroconf", "mock_ble_rpc_device_class") async def test_bluetooth_provision_with_zeroconf_discovery_fast_path( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test zeroconf discovery arrives during WiFi provisioning (fast path - line 551).""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Inject BLE device inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -4272,15 +4308,11 @@ async def test_bluetooth_provision_with_zeroconf_discovery_fast_path( ) # Confirm and scan - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) - # Patch async_provision_wifi to trigger zeroconf discovery - async def mock_provision_wifi(*args, **kwargs): - """Mock provision that triggers zeroconf discovery.""" + # Configure wifi_setconfig to trigger zeroconf discovery + async def mock_wifi_setconfig(*args, **kwargs): + """Mock wifi_setconfig that triggers zeroconf discovery.""" # Trigger zeroconf discovery for the device await hass.config_entries.flow.async_init( DOMAIN, @@ -4297,6 +4329,9 @@ async def test_bluetooth_provision_with_zeroconf_discovery_fast_path( ) # Ensure the zeroconf discovery completes before returning await hass.async_block_till_done() + return {} + + mock_ble_rpc_device.wifi_setconfig.side_effect = mock_wifi_setconfig # Mock device for secure device feature mock_device = create_mock_rpc_device("Test name") @@ -4306,10 +4341,6 @@ async def test_bluetooth_provision_with_zeroconf_discovery_fast_path( "homeassistant.components.shelly.config_flow.PROVISIONING_TIMEOUT", 10, ), - patch( - "homeassistant.components.shelly.config_flow.async_provision_wifi", - side_effect=mock_provision_wifi, - ), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=None, @@ -4342,11 +4373,19 @@ async def test_bluetooth_provision_with_zeroconf_discovery_fast_path( assert len(mock_setup_entry.mock_calls) == 1 -@pytest.mark.usefixtures("mock_zeroconf") +@pytest.mark.usefixtures("mock_zeroconf", "mock_ble_rpc_device_class") async def test_bluetooth_provision_timeout_active_lookup_fails( hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test WiFi provisioning times out and active lookup fails (lines 545-547).""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Configure BLE device to NOT return an IP (so BLE fallback also fails) + mock_ble_rpc_device.status = {"wifi": {"sta_ip": None}} + # Inject BLE device inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -4357,23 +4396,23 @@ async def test_bluetooth_provision_timeout_active_lookup_fails( ) # Confirm and scan - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) # Provision WiFi but no zeroconf discovery arrives, and active lookup fails + # Keep patches active throughout to avoid socket errors from background tasks with ( patch( "homeassistant.components.shelly.config_flow.PROVISIONING_TIMEOUT", 0.01, # Short timeout to trigger timeout path ), - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=None, # Active lookup fails ), + patch( + "homeassistant.components.shelly.config_flow.get_info", + return_value=MOCK_DEVICE_INFO, + ), ): result = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -4387,27 +4426,39 @@ async def test_bluetooth_provision_timeout_active_lookup_fails( # Timeout occurs, active lookup fails, provision unsuccessful result = await hass.config_entries.flow.async_configure(result["flow_id"]) - # Should show provision_failed form - assert result["type"] is FlowResultType.FORM - assert result["step_id"] == "provision_failed" + # Should show provision_failed form + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "provision_failed" - # User aborts after failure - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - side_effect=RuntimeError("BLE device unavailable"), - ): + # User aborts after failure - configure wifi_scan to raise exception + mock_ble_rpc_device.wifi_scan.side_effect = RuntimeError( + "BLE device unavailable" + ) result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) - assert result["type"] is FlowResultType.ABORT - assert result["reason"] == "unknown" + assert result["type"] is FlowResultType.ABORT + assert result["reason"] == "unknown" + + # Allow any pending background tasks to complete + await hass.async_block_till_done(wait_background_tasks=True) +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provision_timeout_ble_fallback_succeeds( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, + mock_ble_rpc_device_class: MagicMock, ) -> None: """Test WiFi provisioning times out, active lookup fails, but BLE fallback succeeds.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Configure BLE device to return IP after provisioning + mock_ble_rpc_device.status = {"wifi": {"sta_ip": "192.168.1.100"}} + # Inject BLE device inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -4417,52 +4468,41 @@ async def test_bluetooth_provision_timeout_ble_fallback_succeeds( context={"source": config_entries.SOURCE_BLUETOOTH}, ) - # Mock device for BLE status query - mock_ble_status_device = AsyncMock() - mock_ble_status_device.status = {"wifi": {"sta_ip": "192.168.1.100"}} + # Mock device for secure device feature (IP connection) + mock_ip_device = AsyncMock() + mock_ip_device.initialize = AsyncMock() + mock_ip_device.name = "Test name" + mock_ip_device.status = {"sys": {}} + mock_ip_device.xmod_info = {} + mock_ip_device.shelly = {"model": MODEL_PLUS_2PM} + mock_ip_device.wifi_setconfig = AsyncMock(return_value={}) + mock_ip_device.ble_setconfig = AsyncMock(return_value={"restart_required": False}) + mock_ip_device.shutdown = AsyncMock() - # Mock device for secure device feature - mock_device = AsyncMock() - mock_device.initialize = AsyncMock() - mock_device.name = "Test name" - mock_device.status = {"sys": {}} - mock_device.xmod_info = {} - mock_device.shelly = {"model": MODEL_PLUS_2PM} - mock_device.wifi_setconfig = AsyncMock(return_value={}) - mock_device.ble_setconfig = AsyncMock(return_value={"restart_required": False}) - mock_device.shutdown = AsyncMock() + # Configure mock_ble_rpc_device_class to return mock_ip_device for IP connections + async def mock_create(aiohttp_session, ws_context, ip_or_options): + if aiohttp_session is None: + return mock_ble_rpc_device + return mock_ip_device + + mock_ble_rpc_device_class.create.side_effect = mock_create # Confirm and scan, then select network and enter password # Provision WiFi but no zeroconf discovery arrives, active lookup fails, BLE fallback succeeds with ( - patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ), patch( "homeassistant.components.shelly.config_flow.PROVISIONING_TIMEOUT", 0.01, # Short timeout to trigger timeout path ), - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=None, # Active lookup fails ), - patch( - "homeassistant.components.shelly.config_flow.ble_rpc_device", - ) as mock_ble_rpc, patch( "homeassistant.components.shelly.config_flow.get_info", return_value=MOCK_DEVICE_INFO, ), - patch( - "homeassistant.components.shelly.config_flow.RpcDevice.create", - return_value=mock_device, - ), ): - # Configure BLE RPC mock to return device with IP - mock_ble_rpc.return_value.__aenter__.return_value = mock_ble_status_device - # Scan for networks result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) @@ -4486,10 +4526,19 @@ async def test_bluetooth_provision_timeout_ble_fallback_succeeds( assert result["data"][CONF_PORT] == DEFAULT_HTTP_PORT +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provision_timeout_ble_fallback_fails( hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test WiFi provisioning times out, active lookup fails, and BLE fallback also fails.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Configure BLE device to return no IP (fallback fails) + mock_ble_rpc_device.status = {"wifi": {"sta_ip": None}} + # Inject BLE device inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -4502,23 +4551,14 @@ async def test_bluetooth_provision_timeout_ble_fallback_fails( # Confirm and scan, select network and enter password # Provision WiFi but no zeroconf discovery, active lookup fails, BLE fallback fails with ( - patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ), patch( "homeassistant.components.shelly.config_flow.PROVISIONING_TIMEOUT", 0.01, # Short timeout to trigger timeout path ), - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=None, # Active lookup fails ), - patch( - "homeassistant.components.shelly.config_flow.async_get_ip_from_ble", - return_value=None, # BLE fallback also fails - ), ): # Scan for networks result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) @@ -4540,21 +4580,25 @@ async def test_bluetooth_provision_timeout_ble_fallback_fails( assert result["type"] is FlowResultType.FORM assert result["step_id"] == "provision_failed" - # User aborts after failure - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - side_effect=RuntimeError("BLE device unavailable"), - ): - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + # User aborts after failure - configure wifi_scan to raise exception + mock_ble_rpc_device.wifi_scan.side_effect = RuntimeError("BLE device unavailable") + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) assert result["type"] is FlowResultType.ABORT assert result["reason"] == "unknown" +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provision_timeout_ble_exception( hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test WiFi provisioning times out, active lookup fails, and BLE raises exception.""" + # Configure mock BLE device for initial wifi scan + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Inject BLE device inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) @@ -4564,30 +4608,29 @@ async def test_bluetooth_provision_timeout_ble_exception( context={"source": config_entries.SOURCE_BLUETOOTH}, ) + # Scan for networks first + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + + # Now configure update_status to raise exception (BLE raises exception during IP fetch) + mock_ble_rpc_device.update_status.side_effect = DeviceConnectionError + # Confirm and scan, select network and enter password # Provision WiFi but no zeroconf discovery, active lookup fails, BLE raises exception + # Keep patches active until all background tasks complete to avoid socket errors with ( - patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ), patch( "homeassistant.components.shelly.config_flow.PROVISIONING_TIMEOUT", 0.01, # Short timeout to trigger timeout path ), - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=None, # Active lookup fails ), patch( - "homeassistant.components.shelly.config_flow.ble_rpc_device", - side_effect=DeviceConnectionError, # BLE raises exception + "homeassistant.components.shelly.config_flow.get_info", + side_effect=DeviceConnectionError, # get_info also fails ), ): - # Scan for networks - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) - # Select network and enter password in single step result = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -4601,27 +4644,36 @@ async def test_bluetooth_provision_timeout_ble_exception( # Timeout occurs, both active lookup and BLE fallback fail (exception) result = await hass.config_entries.flow.async_configure(result["flow_id"]) - # Should show provision_failed form - assert result["type"] is FlowResultType.FORM - assert result["step_id"] == "provision_failed" + # Should show provision_failed form + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "provision_failed" - # User aborts after failure - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - side_effect=RuntimeError("BLE device unavailable"), - ): + # User aborts after failure - configure wifi_scan to raise exception + mock_ble_rpc_device.wifi_scan.side_effect = RuntimeError( + "BLE device unavailable" + ) result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) - assert result["type"] is FlowResultType.ABORT - assert result["reason"] == "unknown" + assert result["type"] is FlowResultType.ABORT + assert result["reason"] == "unknown" + + # Wait for all background tasks to complete before patches are removed + await hass.async_block_till_done(wait_background_tasks=True) +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provision_secure_device_both_enabled( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test provisioning with both AP and BLE disable enabled (default).""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) result = await hass.config_entries.flow.async_init( @@ -4631,14 +4683,10 @@ async def test_bluetooth_provision_secure_device_both_enabled( ) # Confirm with both switches enabled (default) - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {"disable_ap": True, "disable_ble_rpc": True}, - ) + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {"disable_ap": True, "disable_ble_rpc": True}, + ) # Provision and verify security calls mock_device = AsyncMock() @@ -4648,7 +4696,6 @@ async def test_bluetooth_provision_secure_device_both_enabled( mock_device.shutdown = AsyncMock() with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -4678,12 +4725,19 @@ async def test_bluetooth_provision_secure_device_both_enabled( assert mock_device.shutdown.called +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provision_secure_device_both_disabled( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test provisioning with both AP and BLE disable disabled.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) result = await hass.config_entries.flow.async_init( @@ -4693,18 +4747,13 @@ async def test_bluetooth_provision_secure_device_both_disabled( ) # Confirm with both switches disabled - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {"disable_ap": False, "disable_ble_rpc": False}, - ) + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {"disable_ap": False, "disable_ble_rpc": False}, + ) # Provision - with both disabled, secure device method should not create device with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -4725,12 +4774,19 @@ async def test_bluetooth_provision_secure_device_both_disabled( assert result["type"] is FlowResultType.CREATE_ENTRY +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provision_secure_device_only_ap_disabled( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test provisioning with only AP disable enabled.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) result = await hass.config_entries.flow.async_init( @@ -4740,14 +4796,10 @@ async def test_bluetooth_provision_secure_device_only_ap_disabled( ) # Confirm with only AP disable - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {"disable_ap": True, "disable_ble_rpc": False}, - ) + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {"disable_ap": True, "disable_ble_rpc": False}, + ) # Provision and verify only AP disabled mock_device = AsyncMock() @@ -4756,7 +4808,6 @@ async def test_bluetooth_provision_secure_device_only_ap_disabled( mock_device.shutdown = AsyncMock() with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -4785,12 +4836,19 @@ async def test_bluetooth_provision_secure_device_only_ap_disabled( assert mock_device.shutdown.called +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provision_secure_device_only_ble_disabled( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test provisioning with only BLE disable enabled.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) result = await hass.config_entries.flow.async_init( @@ -4800,14 +4858,10 @@ async def test_bluetooth_provision_secure_device_only_ble_disabled( ) # Confirm with only BLE disable - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {"disable_ap": False, "disable_ble_rpc": True}, - ) + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {"disable_ap": False, "disable_ble_rpc": True}, + ) # Provision and verify only BLE disabled mock_device = AsyncMock() @@ -4816,7 +4870,6 @@ async def test_bluetooth_provision_secure_device_only_ble_disabled( mock_device.shutdown = AsyncMock() with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -4845,12 +4898,19 @@ async def test_bluetooth_provision_secure_device_only_ble_disabled( assert mock_device.shutdown.called +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provision_secure_device_with_restart_required( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test provisioning when BLE disable requires restart.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) result = await hass.config_entries.flow.async_init( @@ -4860,14 +4920,10 @@ async def test_bluetooth_provision_secure_device_with_restart_required( ) # Confirm with both enabled - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {"disable_ap": True, "disable_ble_rpc": True}, - ) + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {"disable_ap": True, "disable_ble_rpc": True}, + ) # Provision and verify restart is triggered mock_device = AsyncMock() @@ -4878,7 +4934,6 @@ async def test_bluetooth_provision_secure_device_with_restart_required( mock_device.shutdown = AsyncMock() with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -4907,12 +4962,19 @@ async def test_bluetooth_provision_secure_device_with_restart_required( assert mock_device.shutdown.called +@pytest.mark.usefixtures("mock_ble_rpc_device_class") async def test_bluetooth_provision_secure_device_fails_gracefully( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_setup: AsyncMock, + mock_ble_rpc_device: AsyncMock, ) -> None: """Test provisioning succeeds even when secure device calls fail.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) result = await hass.config_entries.flow.async_init( @@ -4922,14 +4984,10 @@ async def test_bluetooth_provision_secure_device_fails_gracefully( ) # Confirm with both enabled - with patch( - "homeassistant.components.shelly.config_flow.async_scan_wifi_networks", - return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}], - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - {"disable_ap": True, "disable_ble_rpc": True}, - ) + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {"disable_ap": True, "disable_ble_rpc": True}, + ) # Provision with security calls failing - wifi_setconfig will fail mock_device = AsyncMock() @@ -4938,7 +4996,6 @@ async def test_bluetooth_provision_secure_device_fails_gracefully( mock_device.shutdown = AsyncMock() with ( - patch("homeassistant.components.shelly.config_flow.async_provision_wifi"), patch( "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", return_value=("1.1.1.1", 80), @@ -5016,3 +5073,200 @@ async def test_zeroconf_aborts_idle_ble_flow( assert result["result"].unique_id == "C049EF8873E8" assert len(mock_setup.mock_calls) == 1 assert len(mock_setup_entry.mock_calls) == 1 + + +@pytest.mark.usefixtures("mock_zeroconf", "mock_ble_rpc_device_class") +async def test_bluetooth_flow_abort_cleans_up_ble_connection( + hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, +) -> None: + """Test that aborting BLE flow cleans up the BLE connection.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) + + # Start BLE flow + result = await hass.config_entries.flow.async_init( + DOMAIN, + data=BLE_DISCOVERY_INFO, + context={"source": config_entries.SOURCE_BLUETOOTH}, + ) + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "bluetooth_confirm" + flow_id = result["flow_id"] + + # Confirm to proceed to WiFi scan (this creates the BLE connection) + result = await hass.config_entries.flow.async_configure(flow_id, {}) + + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "wifi_scan" + + # Verify BLE device was connected + mock_ble_rpc_device.initialize.assert_called_once() + mock_ble_rpc_device.wifi_scan.assert_called_once() + + # Abort the flow - this should trigger async_remove and clean up BLE + hass.config_entries.flow.async_abort(flow_id) + await hass.async_block_till_done(wait_background_tasks=True) + + # Verify cleanup was called + mock_ble_rpc_device.shutdown.assert_called_once() + + +@pytest.mark.usefixtures("mock_zeroconf") +async def test_bluetooth_ble_initialize_failure_cleans_up( + hass: HomeAssistant, +) -> None: + """Test that initialize failure properly cleans up the device.""" + mock_device = AsyncMock() + mock_device.initialize = AsyncMock(side_effect=DeviceConnectionError) + mock_device.shutdown = AsyncMock() + + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) + + # Start BLE flow + result = await hass.config_entries.flow.async_init( + DOMAIN, + data=BLE_DISCOVERY_INFO, + context={"source": config_entries.SOURCE_BLUETOOTH}, + ) + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "bluetooth_confirm" + + # Confirm to proceed to WiFi scan - this will try to create BLE connection + # but initialize() will fail + with patch( + "homeassistant.components.shelly.config_flow.RpcDevice.create", + return_value=mock_device, + ): + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + + # Should show wifi_scan_failed due to connection error + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "wifi_scan_failed" + + # Verify device was cleaned up after initialize failure + mock_device.shutdown.assert_called_once() + + # Abort the flow to reach terminal state + hass.config_entries.flow.async_abort(result["flow_id"]) + await hass.async_block_till_done(wait_background_tasks=True) + + # Verify shutdown wasn't called again during abort (device was already cleaned up) + mock_device.shutdown.assert_called_once() + + +@pytest.mark.usefixtures("mock_zeroconf", "mock_ble_rpc_device_class") +async def test_bluetooth_ble_shutdown_exception_handled( + hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, +) -> None: + """Test that shutdown exceptions during cleanup are handled gracefully.""" + # Configure mock BLE device + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + # Make shutdown raise an exception + mock_ble_rpc_device.shutdown.side_effect = RuntimeError("Shutdown failed") + + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) + + # Start BLE flow + result = await hass.config_entries.flow.async_init( + DOMAIN, + data=BLE_DISCOVERY_INFO, + context={"source": config_entries.SOURCE_BLUETOOTH}, + ) + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "bluetooth_confirm" + + # Confirm to proceed to WiFi scan (creates BLE connection) + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "wifi_scan" + + # Abort the flow - this should trigger async_remove and cleanup + # The shutdown exception should be caught and logged, not propagate + hass.config_entries.flow.async_abort(result["flow_id"]) + await hass.async_block_till_done(wait_background_tasks=True) + + # Verify shutdown was attempted despite the exception + mock_ble_rpc_device.shutdown.assert_called_once() + + +@pytest.mark.usefixtures("mock_zeroconf", "mock_ble_rpc_device_class") +async def test_bluetooth_provision_ble_reconnect_fails_during_ip_fetch( + hass: HomeAssistant, + mock_ble_rpc_device: AsyncMock, +) -> None: + """Test BLE reconnection fails during IP fetch fallback after provisioning.""" + # Configure mock BLE device for initial wifi scan + mock_ble_rpc_device.wifi_scan.return_value = [ + {"ssid": "MyNetwork", "rssi": -50, "auth": 2} + ] + + # Inject BLE device + inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO) + + result = await hass.config_entries.flow.async_init( + DOMAIN, + data=BLE_DISCOVERY_INFO, + context={"source": config_entries.SOURCE_BLUETOOTH}, + ) + + # Scan for networks first (this creates the initial BLE connection) + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + assert result["step_id"] == "wifi_scan" + + # Track initialize calls - first call succeeds (for wifi_setconfig reconnect), + # second call fails (during _async_get_ip_from_ble reconnect) + init_call_count = 0 + + async def initialize_side_effect() -> None: + nonlocal init_call_count + init_call_count += 1 + if init_call_count == 1: + # First reconnect (for wifi_setconfig) succeeds + return + # Second reconnect (for _async_get_ip_from_ble) fails + raise DeviceConnectionError + + # Simulate: device disconnects, first reconnect succeeds, second fails + mock_ble_rpc_device.connected = False + mock_ble_rpc_device.initialize = AsyncMock(side_effect=initialize_side_effect) + + # Provision WiFi - timeout occurs, active lookup fails, BLE reconnect fails + with ( + patch( + "homeassistant.components.shelly.config_flow.PROVISIONING_TIMEOUT", + 0.01, + ), + patch( + "homeassistant.components.shelly.config_flow.async_lookup_device_by_name", + return_value=None, + ), + patch( + "homeassistant.components.shelly.config_flow.get_info", + side_effect=DeviceConnectionError, + ), + ): + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + {CONF_SSID: "MyNetwork", CONF_PASSWORD: "my_password"}, + ) + + assert result["type"] is FlowResultType.SHOW_PROGRESS + await hass.async_block_till_done() + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + + # Should show provision_failed since BLE reconnection failed during IP fetch + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "provision_failed" + + # Abort flow to reach terminal state + hass.config_entries.flow.async_abort(result["flow_id"]) + await hass.async_block_till_done(wait_background_tasks=True)