Add wifi_scan and extend wifi_setconfig for BLE provisioning (#1014)

Author: bdraco

aioshelly/ble/provisioning.py

@@ -3,10 +3,11 @@
 from __future__ import annotations
 
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING
 
 from ..common import ConnectionOptions
 from ..rpc_device import RpcDevice
+from ..rpc_device.models import ShellyWiFiNetwork
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
@@ -42,7 +43,9 @@ async def ble_rpc_device(ble_device: BLEDevice) -> AsyncIterator[RpcDevice]:
         await device.shutdown()
 
 
-async def async_scan_wifi_networks(ble_device: BLEDevice) -> list[dict[str, Any]]:
+async def async_scan_wifi_networks(
+    ble_device: BLEDevice,
+) -> list[ShellyWiFiNetwork]:
     """Scan for WiFi networks via BLE.
 
     Args:
@@ -57,9 +60,7 @@ async def async_scan_wifi_networks(ble_device: BLEDevice) -> list[dict[str, Any]
 
     """
     async with ble_rpc_device(ble_device) as device:
-        # WiFi scan can take up to 20 seconds - use 30s timeout to be safe
-        scan_result = await device.call_rpc("WiFi.Scan", timeout=30)
-        return cast(list[dict[str, Any]], scan_result.get("results", []))
+        return await device.wifi_scan()
 
 
 async def async_provision_wifi(ble_device: BLEDevice, ssid: str, password: str) -> None:
@@ -76,15 +77,6 @@ async def async_provision_wifi(ble_device: BLEDevice, ssid: str, password: str)
 
     """
     async with ble_rpc_device(ble_device) as device:
-        await device.call_rpc(
-            "WiFi.SetConfig",
-            {
-                "config": {
-                    "sta": {
-                        "ssid": ssid,
-                        "pass": password,
-                        "enable": True,
-                    }
-                }
-            },
+        await device.wifi_setconfig(
+            sta_ssid=ssid, sta_password=password, sta_enable=True
         )

aioshelly/rpc_device/device.py

@@ -47,6 +47,7 @@
     ShellyBLESetConfig,
     ShellyScript,
     ShellyScriptCode,
+    ShellyWiFiNetwork,
     ShellyWiFiSetConfig,
     ShellyWsConfig,
     ShellyWsSetConfig,
@@ -720,12 +721,20 @@ async def ble_getconfig(self) -> ShellyBLEConfig:
         return cast(ShellyBLEConfig, await self.call_rpc("BLE.GetConfig"))
 
     async def wifi_setconfig(
-        self, *, ap_enable: bool | None = None
+        self,
+        *,
+        ap_enable: bool | None = None,
+        sta_ssid: str | None = None,
+        sta_password: str | None = None,
+        sta_enable: bool | None = None,
     ) -> ShellyWiFiSetConfig:
         """Configure WiFi settings with WiFi.SetConfig.
 
         Args:
             ap_enable: Whether to enable the WiFi AP
+            sta_ssid: WiFi station SSID to connect to
+            sta_password: WiFi station password
+            sta_enable: Whether to enable the WiFi station
 
         Returns:
             Response dict, may contain "restart_required": bool
@@ -734,12 +743,39 @@ async def wifi_setconfig(
         config: dict[str, Any] = {}
         if ap_enable is not None:
             config["ap"] = {"enable": ap_enable}
+        if sta_ssid is not None or sta_password is not None or sta_enable is not None:
+            sta_config: dict[str, Any] = {}
+            if sta_ssid is not None:
+                sta_config["ssid"] = sta_ssid
+            if sta_password is not None:
+                sta_config["pass"] = sta_password
+            if sta_enable is not None:
+                sta_config["enable"] = sta_enable
+            config["sta"] = sta_config
 
         return cast(
             ShellyWiFiSetConfig,
             await self.call_rpc("WiFi.SetConfig", {"config": config}),
         )
 
+    async def wifi_scan(self, timeout: float = 30) -> list[ShellyWiFiNetwork]:
+        """Scan for WiFi networks using WiFi.Scan.
+
+        Args:
+            timeout: Scan timeout in seconds (default 30s, scan takes up to 20s)
+
+        Returns:
+            List of WiFi networks with ssid, rssi, auth fields
+
+        Raises:
+            DeviceConnectionError: If connection to device fails
+            RpcCallError: If RPC call fails
+
+        """
+        scan_result = await self.call_rpc("WiFi.Scan", timeout=timeout)
+        results: list[ShellyWiFiNetwork] = scan_result.get("results", [])
+        return results
+
     async def ws_setconfig(
         self, enable: bool, server: str, ssl_ca: str = "*"
     ) -> ShellyWsSetConfig:

aioshelly/rpc_device/models.py

@@ -45,6 +45,16 @@ class ShellyWiFiSetConfig(TypedDict, total=False):
     restart_required: bool
 
 
+class ShellyWiFiNetwork(TypedDict, total=False):
+    """Shelly WiFi Network from scan results."""
+
+    ssid: str
+    bssid: str
+    auth: int
+    channel: int
+    rssi: int
+
+
 class ShellyWsConfig(TypedDict, total=False):
     """Shelly Outbound Websocket Config."""
 

tests/ble/conftest.py

@@ -28,6 +28,8 @@ def mock_rpc_device() -> AsyncMock:
     mock_device = AsyncMock()
     mock_device.initialize = AsyncMock()
     mock_device.call_rpc = AsyncMock()
+    mock_device.wifi_scan = AsyncMock()
+    mock_device.wifi_setconfig = AsyncMock()
     mock_device.shutdown = AsyncMock()
     return mock_device
 

tests/ble/test_provisioning.py

@@ -16,12 +16,10 @@ async def test_scan_wifi_networks_success(
     mock_rpc_device: AsyncMock,
 ) -> None:
     """Test scanning for WiFi networks successfully."""
-    mock_rpc_device.call_rpc.return_value = {
-        "results": [
-            {"ssid": "Network1", "rssi": -50, "auth": 2},
-            {"ssid": "Network2", "rssi": -60, "auth": 3},
-        ]
-    }
+    mock_rpc_device.wifi_scan.return_value = [
+        {"ssid": "Network1", "rssi": -50, "auth": 2},
+        {"ssid": "Network2", "rssi": -60, "auth": 3},
+    ]
 
     result = await async_scan_wifi_networks(mock_ble_device)
 
@@ -30,7 +28,7 @@ async def test_scan_wifi_networks_success(
         {"ssid": "Network2", "rssi": -60, "auth": 3},
     ]
     mock_rpc_device.initialize.assert_called_once()
-    mock_rpc_device.call_rpc.assert_called_once_with("WiFi.Scan", timeout=30)
+    mock_rpc_device.wifi_scan.assert_called_once()
     mock_rpc_device.shutdown.assert_called_once()
 
 
@@ -41,22 +39,7 @@ async def test_scan_wifi_networks_empty_results(
     mock_rpc_device: AsyncMock,
 ) -> None:
     """Test scanning for WiFi networks with empty results."""
-    mock_rpc_device.call_rpc.return_value = {"results": []}
-
-    result = await async_scan_wifi_networks(mock_ble_device)
-
-    assert result == []
-    mock_rpc_device.shutdown.assert_called_once()
-
-
-@pytest.mark.asyncio
-@pytest.mark.usefixtures("mock_rpc_device_class")
-async def test_scan_wifi_networks_no_results_key(
-    mock_ble_device: MagicMock,
-    mock_rpc_device: AsyncMock,
-) -> None:
-    """Test scanning for WiFi networks with missing results key."""
-    mock_rpc_device.call_rpc.return_value = {}
+    mock_rpc_device.wifi_scan.return_value = []
 
     result = await async_scan_wifi_networks(mock_ble_device)
 
@@ -71,7 +54,7 @@ async def test_scan_wifi_networks_exception_cleanup(
     mock_rpc_device: AsyncMock,
 ) -> None:
     """Test that device is shutdown even if scan fails."""
-    mock_rpc_device.call_rpc.side_effect = Exception("Scan failed")
+    mock_rpc_device.wifi_scan.side_effect = Exception("Scan failed")
 
     with pytest.raises(Exception, match="Scan failed"):
         await async_scan_wifi_networks(mock_ble_device)
@@ -90,17 +73,10 @@ async def test_provision_wifi_success(
     await async_provision_wifi(mock_ble_device, "MyNetwork", "MyPassword")
 
     mock_rpc_device.initialize.assert_called_once()
-    mock_rpc_device.call_rpc.assert_called_once_with(
-        "WiFi.SetConfig",
-        {
-            "config": {
-                "sta": {
-                    "ssid": "MyNetwork",
-                    "pass": "MyPassword",
-                    "enable": True,
-                }
-            }
-        },
+    mock_rpc_device.wifi_setconfig.assert_called_once_with(
+        sta_ssid="MyNetwork",
+        sta_password="MyPassword",  # noqa: S106
+        sta_enable=True,
     )
     mock_rpc_device.shutdown.assert_called_once()
 
@@ -112,7 +88,7 @@ async def test_provision_wifi_exception_cleanup(
     mock_rpc_device: AsyncMock,
 ) -> None:
     """Test that device is shutdown even if provisioning fails."""
-    mock_rpc_device.call_rpc.side_effect = Exception("Provisioning failed")
+    mock_rpc_device.wifi_setconfig.side_effect = Exception("Provisioning failed")
 
     with pytest.raises(Exception, match="Provisioning failed"):
         await async_provision_wifi(mock_ble_device, "MyNetwork", "MyPassword")

tests/rpc_device/test_device.py

@@ -549,6 +549,94 @@ async def test_wifi_setconfig(rpc_device: RpcDevice) -> None:
     }
 
 
+@pytest.mark.asyncio
+async def test_wifi_setconfig_sta(rpc_device: RpcDevice) -> None:
+    """Test RpcDevice wifi_setconfig method with STA configuration."""
+    await rpc_device.wifi_setconfig(
+        sta_ssid="MyNetwork",
+        sta_password="MyPassword",  # noqa: S106
+        sta_enable=True,
+    )
+
+    assert rpc_device.call_rpc_multiple.call_count == 1
+    assert rpc_device.call_rpc_multiple.call_args[0][0][0][0] == "WiFi.SetConfig"
+    assert rpc_device.call_rpc_multiple.call_args[0][0][0][1] == {
+        "config": {
+            "sta": {
+                "ssid": "MyNetwork",
+                "pass": "MyPassword",
+                "enable": True,
+            }
+        }
+    }
+
+
+@pytest.mark.asyncio
+async def test_wifi_setconfig_ap_and_sta(rpc_device: RpcDevice) -> None:
+    """Test RpcDevice wifi_setconfig method with both AP and STA configuration."""
+    await rpc_device.wifi_setconfig(
+        ap_enable=False,
+        sta_ssid="MyNetwork",
+        sta_password="MyPassword",  # noqa: S106
+        sta_enable=True,
+    )
+
+    assert rpc_device.call_rpc_multiple.call_count == 1
+    assert rpc_device.call_rpc_multiple.call_args[0][0][0][0] == "WiFi.SetConfig"
+    assert rpc_device.call_rpc_multiple.call_args[0][0][0][1] == {
+        "config": {
+            "ap": {"enable": False},
+            "sta": {
+                "ssid": "MyNetwork",
+                "pass": "MyPassword",
+                "enable": True,
+            },
+        }
+    }
+
+
+@pytest.mark.asyncio
+async def test_wifi_scan(rpc_device: RpcDevice) -> None:
+    """Test RpcDevice wifi_scan method."""
+    rpc_device.call_rpc_multiple.return_value = [
+        {
+            "results": [
+                {"ssid": "Network1", "rssi": -50, "auth": 2},
+                {"ssid": "Network2", "rssi": -60, "auth": 3},
+            ]
+        }
+    ]
+
+    result = await rpc_device.wifi_scan()
+
+    assert result == [
+        {"ssid": "Network1", "rssi": -50, "auth": 2},
+        {"ssid": "Network2", "rssi": -60, "auth": 3},
+    ]
+    assert rpc_device.call_rpc_multiple.call_count == 1
+    assert rpc_device.call_rpc_multiple.call_args[0][0][0][0] == "WiFi.Scan"
+
+
+@pytest.mark.asyncio
+async def test_wifi_scan_empty(rpc_device: RpcDevice) -> None:
+    """Test RpcDevice wifi_scan method with empty results."""
+    rpc_device.call_rpc_multiple.return_value = [{"results": []}]
+
+    result = await rpc_device.wifi_scan()
+
+    assert result == []
+
+
+@pytest.mark.asyncio
+async def test_wifi_scan_no_results_key(rpc_device: RpcDevice) -> None:
+    """Test RpcDevice wifi_scan method with missing results key."""
+    rpc_device.call_rpc_multiple.return_value = [{}]
+
+    result = await rpc_device.wifi_scan()
+
+    assert result == []
+
+
 @pytest.mark.asyncio
 async def test_script_stop(rpc_device: RpcDevice) -> None:
     """Test RpcDevice script_stop method."""