Add a channel changing API to ZHA (#92076)

* Expose channel changing over the websocket API

* Expose channel changing as a service

* Type annotate some existing unit test fixtures

* Add unit tests

* Rename `api.change_channel` to `api.async_change_channel`

* Expand on channel migration in the service description

* Remove channel changing service, we only really need the websocket API

* Update homeassistant/components/zha/websocket_api.py

* Black

---------

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>
Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/components/zha/api.py

@@ -2,10 +2,12 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 from zigpy.backups import NetworkBackup
 from zigpy.config import CONF_DEVICE, CONF_DEVICE_PATH
+from zigpy.types import Channels
+from zigpy.util import pick_optimal_channel
 
 from .core.const import (
     CONF_RADIO_TYPE,
@@ -111,3 +113,22 @@ def async_get_radio_path(
         config_entry = _get_config_entry(hass)
 
     return config_entry.data[CONF_DEVICE][CONF_DEVICE_PATH]
+
+
+async def async_change_channel(
+    hass: HomeAssistant, new_channel: int | Literal["auto"]
+) -> None:
+    """Migrate the ZHA network to a new channel."""
+
+    zha_gateway: ZHAGateway = _get_gateway(hass)
+    app = zha_gateway.application_controller
+
+    if new_channel == "auto":
+        channel_energy = await app.energy_scan(
+            channels=Channels.ALL_CHANNELS,
+            duration_exp=4,
+            count=1,
+        )
+        new_channel = pick_optimal_channel(channel_energy)
+
+    await app.move_network_to_channel(new_channel)

homeassistant/components/zha/websocket_api.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 import asyncio
 import logging
-from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypeVar, cast
 
 import voluptuous as vol
 import zigpy.backups
@@ -19,7 +19,11 @@ import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.service import async_register_admin_service
 
-from .api import async_get_active_network_settings, async_get_radio_type
+from .api import (
+    async_change_channel,
+    async_get_active_network_settings,
+    async_get_radio_type,
+)
 from .core.const import (
     ATTR_ARGS,
     ATTR_ATTRIBUTE,
@@ -93,6 +97,7 @@ ATTR_DURATION = "duration"
 ATTR_GROUP = "group"
 ATTR_IEEE_ADDRESS = "ieee_address"
 ATTR_INSTALL_CODE = "install_code"
+ATTR_NEW_CHANNEL = "new_channel"
 ATTR_SOURCE_IEEE = "source_ieee"
 ATTR_TARGET_IEEE = "target_ieee"
 ATTR_QR_CODE = "qr_code"
@@ -1204,6 +1209,23 @@ async def websocket_restore_network_backup(
         connection.send_result(msg[ID])
 
 
+@websocket_api.require_admin
+@websocket_api.websocket_command(
+    {
+        vol.Required(TYPE): "zha/network/change_channel",
+        vol.Required(ATTR_NEW_CHANNEL): vol.Any("auto", vol.Range(11, 26)),
+    }
+)
+@websocket_api.async_response
+async def websocket_change_channel(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Migrate the Zigbee network to a new channel."""
+    new_channel = cast(Literal["auto"] | int, msg[ATTR_NEW_CHANNEL])
+    await async_change_channel(hass, new_channel=new_channel)
+    connection.send_result(msg[ID])
+
+
 @callback
 def async_load_api(hass: HomeAssistant) -> None:
     """Set up the web socket API."""
@@ -1527,6 +1549,7 @@ def async_load_api(hass: HomeAssistant) -> None:
     websocket_api.async_register_command(hass, websocket_list_network_backups)
     websocket_api.async_register_command(hass, websocket_create_network_backup)
     websocket_api.async_register_command(hass, websocket_restore_network_backup)
+    websocket_api.async_register_command(hass, websocket_change_channel)
 
 
 @callback

tests/components/zha/conftest.py

@@ -116,7 +116,9 @@ def zigpy_app_controller():
     app.state.network_info.channel = 15
     app.state.network_info.network_key.key = zigpy.types.KeyData(range(16))
 
-    with patch("zigpy.device.Device.request"):
+    with patch("zigpy.device.Device.request"), patch.object(
+        app, "permit", autospec=True
+    ), patch.object(app, "permit_with_key", autospec=True):
         yield app
 
 

tests/components/zha/test_api.py

@@ -1,5 +1,8 @@
 """Test ZHA API."""
-from unittest.mock import patch
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from unittest.mock import call, patch
 
 import pytest
 import zigpy.backups
@@ -10,6 +13,9 @@ from homeassistant.components.zha import api
 from homeassistant.components.zha.core.const import RadioType
 from homeassistant.core import HomeAssistant
 
+if TYPE_CHECKING:
+    from zigpy.application import ControllerApplication
+
 
 @pytest.fixture(autouse=True)
 def required_platform_only():
@@ -29,7 +35,7 @@ async def test_async_get_network_settings_active(
 
 
 async def test_async_get_network_settings_inactive(
-    hass: HomeAssistant, setup_zha, zigpy_app_controller
+    hass: HomeAssistant, setup_zha, zigpy_app_controller: ControllerApplication
 ) -> None:
     """Test reading settings with an inactive ZHA installation."""
     await setup_zha()
@@ -59,7 +65,7 @@ async def test_async_get_network_settings_inactive(
 
 
 async def test_async_get_network_settings_missing(
-    hass: HomeAssistant, setup_zha, zigpy_app_controller
+    hass: HomeAssistant, setup_zha, zigpy_app_controller: ControllerApplication
 ) -> None:
     """Test reading settings with an inactive ZHA installation, no valid channel."""
     await setup_zha()
@@ -100,3 +106,38 @@ async def test_async_get_radio_path_active(hass: HomeAssistant, setup_zha) -> No
 
     radio_path = api.async_get_radio_path(hass)
     assert radio_path == "/dev/ttyUSB0"
+
+
+async def test_change_channel(
+    hass: HomeAssistant, setup_zha, zigpy_app_controller: ControllerApplication
+) -> None:
+    """Test changing the channel."""
+    await setup_zha()
+
+    with patch.object(
+        zigpy_app_controller, "move_network_to_channel", autospec=True
+    ) as mock_move_network_to_channel:
+        await api.async_change_channel(hass, 20)
+
+    assert mock_move_network_to_channel.mock_calls == [call(20)]
+
+
+async def test_change_channel_auto(
+    hass: HomeAssistant, setup_zha, zigpy_app_controller: ControllerApplication
+) -> None:
+    """Test changing the channel automatically using an energy scan."""
+    await setup_zha()
+
+    with patch.object(
+        zigpy_app_controller, "move_network_to_channel", autospec=True
+    ) as mock_move_network_to_channel, patch.object(
+        zigpy_app_controller,
+        "energy_scan",
+        autospec=True,
+        return_value={c: c for c in range(11, 26 + 1)},
+    ), patch.object(
+        api, "pick_optimal_channel", autospec=True, return_value=25
+    ):
+        await api.async_change_channel(hass, "auto")
+
+    assert mock_move_network_to_channel.mock_calls == [call(25)]

tests/components/zha/test_websocket_api.py

@@ -1,7 +1,10 @@
 """Test ZHA WebSocket API."""
+from __future__ import annotations
+
 from binascii import unhexlify
 from copy import deepcopy
-from unittest.mock import AsyncMock, patch
+from typing import TYPE_CHECKING
+from unittest.mock import ANY, AsyncMock, call, patch
 
 import pytest
 import voluptuous as vol
@@ -24,8 +27,6 @@ from homeassistant.components.zha.core.const import (
     ATTR_NEIGHBORS,
     ATTR_QUIRK_APPLIED,
     CLUSTER_TYPE_IN,
-    DATA_ZHA,
-    DATA_ZHA_GATEWAY,
     EZSP_OVERWRITE_EUI64,
     GROUP_ID,
     GROUP_IDS,
@@ -59,6 +60,9 @@ from tests.common import MockUser
 IEEE_SWITCH_DEVICE = "01:2d:6f:00:0a:90:69:e7"
 IEEE_GROUPABLE_DEVICE = "01:2d:6f:00:0a:90:69:e8"
 
+if TYPE_CHECKING:
+    from zigpy.application import ControllerApplication
+
 
 @pytest.fixture(autouse=True)
 def required_platform_only():
@@ -282,15 +286,17 @@ async def test_get_zha_config_with_alarm(
     assert configuration == BASE_CUSTOM_CONFIGURATION
 
 
-async def test_update_zha_config(zha_client, zigpy_app_controller) -> None:
+async def test_update_zha_config(
+    zha_client, app_controller: ControllerApplication
+) -> None:
     """Test updating ZHA custom configuration."""
 
-    configuration = deepcopy(CONFIG_WITH_ALARM_OPTIONS)
+    configuration: dict = deepcopy(CONFIG_WITH_ALARM_OPTIONS)
     configuration["data"]["zha_options"]["default_light_transition"] = 10
 
     with patch(
         "bellows.zigbee.application.ControllerApplication.new",
-        return_value=zigpy_app_controller,
+        return_value=app_controller,
     ):
         await zha_client.send_json(
             {ID: 5, TYPE: "zha/configuration/update", "data": configuration["data"]}
@@ -463,14 +469,12 @@ async def test_remove_group(zha_client) -> None:
 
 
 @pytest.fixture
-async def app_controller(hass, setup_zha):
+async def app_controller(
+    hass: HomeAssistant, setup_zha, zigpy_app_controller: ControllerApplication
+) -> ControllerApplication:
     """Fixture for zigpy Application Controller."""
     await setup_zha()
-    controller = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].application_controller
-    p1 = patch.object(controller, "permit")
-    p2 = patch.object(controller, "permit_with_key", new=AsyncMock())
-    with p1, p2:
-        yield controller
+    return zigpy_app_controller
 
 
 @pytest.mark.parametrize(
@@ -492,7 +496,7 @@ async def app_controller(hass, setup_zha):
 )
 async def test_permit_ha12(
     hass: HomeAssistant,
-    app_controller,
+    app_controller: ControllerApplication,
     hass_admin_user: MockUser,
     params,
     duration,
@@ -532,7 +536,7 @@ IC_TEST_PARAMS = (
 @pytest.mark.parametrize(("params", "src_ieee", "code"), IC_TEST_PARAMS)
 async def test_permit_with_install_code(
     hass: HomeAssistant,
-    app_controller,
+    app_controller: ControllerApplication,
     hass_admin_user: MockUser,
     params,
     src_ieee,
@@ -587,7 +591,10 @@ IC_FAIL_PARAMS = (
 
 @pytest.mark.parametrize("params", IC_FAIL_PARAMS)
 async def test_permit_with_install_code_fail(
-    hass: HomeAssistant, app_controller, hass_admin_user: MockUser, params
+    hass: HomeAssistant,
+    app_controller: ControllerApplication,
+    hass_admin_user: MockUser,
+    params,
 ) -> None:
     """Test permit service with install code."""
 
@@ -626,7 +633,7 @@ IC_QR_CODE_TEST_PARAMS = (
 @pytest.mark.parametrize(("params", "src_ieee", "code"), IC_QR_CODE_TEST_PARAMS)
 async def test_permit_with_qr_code(
     hass: HomeAssistant,
-    app_controller,
+    app_controller: ControllerApplication,
     hass_admin_user: MockUser,
     params,
     src_ieee,
@@ -646,7 +653,7 @@ async def test_permit_with_qr_code(
 
 @pytest.mark.parametrize(("params", "src_ieee", "code"), IC_QR_CODE_TEST_PARAMS)
 async def test_ws_permit_with_qr_code(
-    app_controller, zha_client, params, src_ieee, code
+    app_controller: ControllerApplication, zha_client, params, src_ieee, code
 ) -> None:
     """Test permit service with install code from qr code."""
 
@@ -668,7 +675,7 @@ async def test_ws_permit_with_qr_code(
 
 @pytest.mark.parametrize("params", IC_FAIL_PARAMS)
 async def test_ws_permit_with_install_code_fail(
-    app_controller, zha_client, params
+    app_controller: ControllerApplication, zha_client, params
 ) -> None:
     """Test permit ws service with install code."""
 
@@ -703,7 +710,7 @@ async def test_ws_permit_with_install_code_fail(
     ),
 )
 async def test_ws_permit_ha12(
-    app_controller, zha_client, params, duration, node
+    app_controller: ControllerApplication, zha_client, params, duration, node
 ) -> None:
     """Test permit ws service."""
 
@@ -722,7 +729,9 @@ async def test_ws_permit_ha12(
     assert app_controller.permit_with_key.call_count == 0
 
 
-async def test_get_network_settings(app_controller, zha_client) -> None:
+async def test_get_network_settings(
+    app_controller: ControllerApplication, zha_client
+) -> None:
     """Test current network settings are returned."""
 
     await app_controller.backups.create_backup()
@@ -737,7 +746,9 @@ async def test_get_network_settings(app_controller, zha_client) -> None:
     assert "network_info" in msg["result"]["settings"]
 
 
-async def test_list_network_backups(app_controller, zha_client) -> None:
+async def test_list_network_backups(
+    app_controller: ControllerApplication, zha_client
+) -> None:
     """Test backups are serialized."""
 
     await app_controller.backups.create_backup()
@@ -751,7 +762,9 @@ async def test_list_network_backups(app_controller, zha_client) -> None:
     assert "network_info" in msg["result"][0]
 
 
-async def test_create_network_backup(app_controller, zha_client) -> None:
+async def test_create_network_backup(
+    app_controller: ControllerApplication, zha_client
+) -> None:
     """Test creating backup."""
 
     assert not app_controller.backups.backups
@@ -765,7 +778,9 @@ async def test_create_network_backup(app_controller, zha_client) -> None:
     assert "backup" in msg["result"] and "is_complete" in msg["result"]
 
 
-async def test_restore_network_backup_success(app_controller, zha_client) -> None:
+async def test_restore_network_backup_success(
+    app_controller: ControllerApplication, zha_client
+) -> None:
     """Test successfully restoring a backup."""
 
     backup = zigpy.backups.NetworkBackup()
@@ -789,7 +804,7 @@ async def test_restore_network_backup_success(app_controller, zha_client) -> Non
 
 
 async def test_restore_network_backup_force_write_eui64(
-    app_controller, zha_client
+    app_controller: ControllerApplication, zha_client
 ) -> None:
     """Test successfully restoring a backup."""
 
@@ -821,7 +836,9 @@ async def test_restore_network_backup_force_write_eui64(
 
 
 @patch("zigpy.backups.NetworkBackup.from_dict", new=lambda v: v)
-async def test_restore_network_backup_failure(app_controller, zha_client) -> None:
+async def test_restore_network_backup_failure(
+    app_controller: ControllerApplication, zha_client
+) -> None:
     """Test successfully restoring a backup."""
 
     with patch.object(
@@ -840,3 +857,29 @@ async def test_restore_network_backup_failure(app_controller, zha_client) -> Non
     assert msg["type"] == const.TYPE_RESULT
     assert not msg["success"]
     assert msg["error"]["code"] == const.ERR_INVALID_FORMAT
+
+
+@pytest.mark.parametrize("new_channel", ["auto", 15])
+async def test_websocket_change_channel(
+    new_channel: int | str, app_controller: ControllerApplication, zha_client
+) -> None:
+    """Test websocket API to migrate the network to a new channel."""
+
+    with patch(
+        "homeassistant.components.zha.websocket_api.async_change_channel",
+        autospec=True,
+    ) as change_channel_mock:
+        await zha_client.send_json(
+            {
+                ID: 6,
+                TYPE: f"{DOMAIN}/network/change_channel",
+                "new_channel": new_channel,
+            }
+        )
+        msg = await zha_client.receive_json()
+
+    assert msg["id"] == 6
+    assert msg["type"] == const.TYPE_RESULT
+    assert msg["success"]
+
+    change_channel_mock.mock_calls == [call(ANY, new_channel)]