From f9ac1f38390306b2bd37b87ce008cb1604c5844b Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Thu, 27 Apr 2023 11:04:22 -0400 Subject: [PATCH] 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 Co-authored-by: Paulus Schoutsen --- homeassistant/components/zha/api.py | 23 ++++- homeassistant/components/zha/websocket_api.py | 27 +++++- tests/components/zha/conftest.py | 4 +- tests/components/zha/test_api.py | 47 +++++++++- tests/components/zha/test_websocket_api.py | 93 ++++++++++++++----- 5 files changed, 162 insertions(+), 32 deletions(-) diff --git a/homeassistant/components/zha/api.py b/homeassistant/components/zha/api.py index 652f19d24ba..3d44103e225 100644 --- a/homeassistant/components/zha/api.py +++ b/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) diff --git a/homeassistant/components/zha/websocket_api.py b/homeassistant/components/zha/websocket_api.py index 322107a074e..2d4126861b4 100644 --- a/homeassistant/components/zha/websocket_api.py +++ b/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 diff --git a/tests/components/zha/conftest.py b/tests/components/zha/conftest.py index 8e31b45afd8..0621c6521a9 100644 --- a/tests/components/zha/conftest.py +++ b/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 diff --git a/tests/components/zha/test_api.py b/tests/components/zha/test_api.py index b75a65ed5b6..85f85cc0437 100644 --- a/tests/components/zha/test_api.py +++ b/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)] diff --git a/tests/components/zha/test_websocket_api.py b/tests/components/zha/test_websocket_api.py index 7a24daaa3ba..720cfaaac9b 100644 --- a/tests/components/zha/test_websocket_api.py +++ b/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)]