From 6e22251e1d3a4d9a01fea59cd4e4021fbe0226a6 Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Tue, 20 Apr 2021 21:40:54 -0400 Subject: [PATCH] Add support to enable/disable zwave_js data collection (#49440) --- homeassistant/components/zwave_js/__init__.py | 10 +- homeassistant/components/zwave_js/api.py | 163 +++++++++++++++--- homeassistant/components/zwave_js/const.py | 1 + homeassistant/components/zwave_js/helpers.py | 18 +- tests/components/zwave_js/test_api.py | 75 +++++++- tests/components/zwave_js/test_init.py | 48 ++++++ 6 files changed, 287 insertions(+), 28 deletions(-) diff --git a/homeassistant/components/zwave_js/__init__.py b/homeassistant/components/zwave_js/__init__.py index 37d85b81ebe..b7d95ab7bc7 100644 --- a/homeassistant/components/zwave_js/__init__.py +++ b/homeassistant/components/zwave_js/__init__.py @@ -50,6 +50,7 @@ from .const import ( ATTR_TYPE, ATTR_VALUE, ATTR_VALUE_RAW, + CONF_DATA_COLLECTION_OPTED_IN, CONF_INTEGRATION_CREATED_ADDON, CONF_NETWORK_KEY, CONF_USB_PATH, @@ -64,7 +65,7 @@ from .const import ( ZWAVE_JS_VALUE_NOTIFICATION_EVENT, ) from .discovery import async_discover_values -from .helpers import get_device_id +from .helpers import async_enable_statistics, get_device_id from .migrate import async_migrate_discovered_value from .services import ZWaveServices @@ -322,6 +323,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: LOGGER.info("Connection to Zwave JS Server initialized") + # If opt in preference hasn't been specified yet, we do nothing, otherwise + # we apply the preference + if opted_in := entry.data.get(CONF_DATA_COLLECTION_OPTED_IN): + await async_enable_statistics(client) + elif opted_in is False: + await client.driver.async_disable_statistics() + # Check for nodes that no longer exist and remove them stored_devices = device_registry.async_entries_for_config_entry( dev_reg, entry.entry_id diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py index 2792fc6819b..fc4f16bda33 100644 --- a/homeassistant/components/zwave_js/api.py +++ b/homeassistant/components/zwave_js/api.py @@ -2,11 +2,14 @@ from __future__ import annotations import dataclasses +from functools import wraps import json +from typing import Callable from aiohttp import hdrs, web, web_exceptions import voluptuous as vol from zwave_js_server import dump +from zwave_js_server.client import Client from zwave_js_server.const import LogLevel from zwave_js_server.exceptions import InvalidNewValue, NotFoundError, SetValueFailed from zwave_js_server.model.log_config import LogConfig @@ -20,6 +23,7 @@ from homeassistant.components.websocket_api.const import ( ERR_NOT_SUPPORTED, ERR_UNKNOWN_ERROR, ) +from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_URL from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import config_validation as cv @@ -27,7 +31,13 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.device_registry import DeviceEntry from homeassistant.helpers.dispatcher import async_dispatcher_connect -from .const import DATA_CLIENT, DOMAIN, EVENT_DEVICE_ADDED_TO_REGISTRY +from .const import ( + CONF_DATA_COLLECTION_OPTED_IN, + DATA_CLIENT, + DOMAIN, + EVENT_DEVICE_ADDED_TO_REGISTRY, +) +from .helpers import async_enable_statistics, update_data_collection_preference # general API constants ID = "id" @@ -50,6 +60,26 @@ FORCE_CONSOLE = "force_console" VALUE_ID = "value_id" STATUS = "status" +# constants for data collection +ENABLED = "enabled" +OPTED_IN = "opted_in" + + +def async_get_entry(orig_func: Callable) -> Callable: + """Decorate async function to get entry.""" + + @wraps(orig_func) + async def async_get_entry_func( + hass: HomeAssistant, connection: ActiveConnection, msg: dict + ) -> None: + """Provide user specific data and store to function.""" + entry_id = msg[ENTRY_ID] + entry = hass.config_entries.async_get_entry(entry_id) + client = hass.data[DOMAIN][entry_id][DATA_CLIENT] + await orig_func(hass, connection, msg, entry, client) + + return async_get_entry_func + @callback def async_register_api(hass: HomeAssistant) -> None: @@ -65,6 +95,10 @@ def async_register_api(hass: HomeAssistant) -> None: websocket_api.async_register_command(hass, websocket_get_log_config) websocket_api.async_register_command(hass, websocket_get_config_parameters) websocket_api.async_register_command(hass, websocket_set_config_parameter) + websocket_api.async_register_command( + hass, websocket_update_data_collection_preference + ) + websocket_api.async_register_command(hass, websocket_data_collection_status) hass.http.register_view(DumpView) # type: ignore @@ -140,12 +174,15 @@ def websocket_node_status( vol.Optional("secure", default=False): bool, } ) +@async_get_entry async def websocket_add_node( - hass: HomeAssistant, connection: ActiveConnection, msg: dict + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, ) -> None: """Add a node to the Z-Wave network.""" - entry_id = msg[ENTRY_ID] - client = hass.data[DOMAIN][entry_id][DATA_CLIENT] controller = client.driver.controller include_non_secure = not msg["secure"] @@ -210,12 +247,15 @@ async def websocket_add_node( vol.Required(ENTRY_ID): str, } ) +@async_get_entry async def websocket_stop_inclusion( - hass: HomeAssistant, connection: ActiveConnection, msg: dict + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, ) -> None: """Cancel adding a node to the Z-Wave network.""" - entry_id = msg[ENTRY_ID] - client = hass.data[DOMAIN][entry_id][DATA_CLIENT] controller = client.driver.controller result = await controller.async_stop_inclusion() connection.send_result( @@ -232,12 +272,15 @@ async def websocket_stop_inclusion( vol.Required(ENTRY_ID): str, } ) +@async_get_entry async def websocket_stop_exclusion( - hass: HomeAssistant, connection: ActiveConnection, msg: dict + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, ) -> None: """Cancel removing a node from the Z-Wave network.""" - entry_id = msg[ENTRY_ID] - client = hass.data[DOMAIN][entry_id][DATA_CLIENT] controller = client.driver.controller result = await controller.async_stop_exclusion() connection.send_result( @@ -254,12 +297,15 @@ async def websocket_stop_exclusion( vol.Required(ENTRY_ID): str, } ) +@async_get_entry async def websocket_remove_node( - hass: HomeAssistant, connection: ActiveConnection, msg: dict + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, ) -> None: """Remove a node from the Z-Wave network.""" - entry_id = msg[ENTRY_ID] - client = hass.data[DOMAIN][entry_id][DATA_CLIENT] controller = client.driver.controller @callback @@ -311,13 +357,16 @@ async def websocket_remove_node( vol.Required(NODE_ID): int, }, ) +@async_get_entry async def websocket_refresh_node_info( - hass: HomeAssistant, connection: ActiveConnection, msg: dict + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, ) -> None: """Re-interview a node.""" - entry_id = msg[ENTRY_ID] node_id = msg[NODE_ID] - client = hass.data[DOMAIN][entry_id][DATA_CLIENT] node = client.driver.controller.nodes.get(node_id) if node is None: @@ -340,16 +389,19 @@ async def websocket_refresh_node_info( vol.Required(VALUE): int, } ) +@async_get_entry async def websocket_set_config_parameter( - hass: HomeAssistant, connection: ActiveConnection, msg: dict + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, ) -> None: """Set a config parameter value for a Z-Wave node.""" - entry_id = msg[ENTRY_ID] node_id = msg[NODE_ID] property_ = msg[PROPERTY] property_key = msg.get(PROPERTY_KEY) value = msg[VALUE] - client = hass.data[DOMAIN][entry_id][DATA_CLIENT] node = client.driver.controller.nodes[node_id] try: zwave_value, cmd_status = await async_set_config_parameter( @@ -464,12 +516,15 @@ def filename_is_present_if_logging_to_file(obj: dict) -> dict: ), }, ) +@async_get_entry async def websocket_update_log_config( - hass: HomeAssistant, connection: ActiveConnection, msg: dict + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, ) -> None: """Update the driver log config.""" - entry_id = msg[ENTRY_ID] - client = hass.data[DOMAIN][entry_id][DATA_CLIENT] await client.driver.async_update_log_config(LogConfig(**msg[CONFIG])) connection.send_result( msg[ID], @@ -484,12 +539,15 @@ async def websocket_update_log_config( vol.Required(ENTRY_ID): str, }, ) +@async_get_entry async def websocket_get_log_config( - hass: HomeAssistant, connection: ActiveConnection, msg: dict + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, ) -> None: """Get log configuration for the Z-Wave JS driver.""" - entry_id = msg[ENTRY_ID] - client = hass.data[DOMAIN][entry_id][DATA_CLIENT] result = await client.driver.async_get_log_config() connection.send_result( msg[ID], @@ -497,6 +555,61 @@ async def websocket_get_log_config( ) +@websocket_api.require_admin # type: ignore +@websocket_api.async_response +@websocket_api.websocket_command( + { + vol.Required(TYPE): "zwave_js/update_data_collection_preference", + vol.Required(ENTRY_ID): str, + vol.Required(OPTED_IN): bool, + }, +) +@async_get_entry +async def websocket_update_data_collection_preference( + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, +) -> None: + """Update preference for data collection and enable/disable collection.""" + opted_in = msg[OPTED_IN] + update_data_collection_preference(hass, entry, opted_in) + + if opted_in: + await async_enable_statistics(client) + else: + await client.driver.async_disable_statistics() + + connection.send_result( + msg[ID], + ) + + +@websocket_api.require_admin # type: ignore +@websocket_api.async_response +@websocket_api.websocket_command( + { + vol.Required(TYPE): "zwave_js/data_collection_status", + vol.Required(ENTRY_ID): str, + }, +) +@async_get_entry +async def websocket_data_collection_status( + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict, + entry: ConfigEntry, + client: Client, +) -> None: + """Return data collection preference and status.""" + result = { + OPTED_IN: entry.data.get(CONF_DATA_COLLECTION_OPTED_IN), + ENABLED: await client.driver.async_is_statistics_enabled(), + } + connection.send_result(msg[ID], result) + + class DumpView(HomeAssistantView): """View to dump the state of the Z-Wave JS server.""" diff --git a/homeassistant/components/zwave_js/const.py b/homeassistant/components/zwave_js/const.py index afd899e0ee0..54c1ca78e30 100644 --- a/homeassistant/components/zwave_js/const.py +++ b/homeassistant/components/zwave_js/const.py @@ -7,6 +7,7 @@ CONF_INTEGRATION_CREATED_ADDON = "integration_created_addon" CONF_NETWORK_KEY = "network_key" CONF_USB_PATH = "usb_path" CONF_USE_ADDON = "use_addon" +CONF_DATA_COLLECTION_OPTED_IN = "data_collection_opted_in" DOMAIN = "zwave_js" DATA_CLIENT = "client" diff --git a/homeassistant/components/zwave_js/helpers.py b/homeassistant/components/zwave_js/helpers.py index d535a22394c..98c308ea58c 100644 --- a/homeassistant/components/zwave_js/helpers.py +++ b/homeassistant/components/zwave_js/helpers.py @@ -7,11 +7,27 @@ from zwave_js_server.client import Client as ZwaveClient from zwave_js_server.model.node import Node as ZwaveNode from homeassistant.config_entries import ConfigEntry +from homeassistant.const import __version__ as HA_VERSION from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import async_get as async_get_dev_reg from homeassistant.helpers.entity_registry import async_get as async_get_ent_reg -from .const import DATA_CLIENT, DOMAIN +from .const import CONF_DATA_COLLECTION_OPTED_IN, DATA_CLIENT, DOMAIN + + +async def async_enable_statistics(client: ZwaveClient) -> None: + """Enable statistics on the driver.""" + await client.driver.async_enable_statistics("Home Assistant", HA_VERSION) + + +@callback +def update_data_collection_preference( + hass: HomeAssistant, entry: ConfigEntry, preference: bool +) -> None: + """Update data collection preference on config entry.""" + new_data = entry.data.copy() + new_data[CONF_DATA_COLLECTION_OPTED_IN] = preference + hass.config_entries.async_update_entry(entry, data=new_data) @callback diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py index ee718020b7a..edd711b07d1 100644 --- a/tests/components/zwave_js/test_api.py +++ b/tests/components/zwave_js/test_api.py @@ -17,12 +17,16 @@ from homeassistant.components.zwave_js.api import ( LEVEL, LOG_TO_FILE, NODE_ID, + OPTED_IN, PROPERTY, PROPERTY_KEY, TYPE, VALUE, ) -from homeassistant.components.zwave_js.const import DOMAIN +from homeassistant.components.zwave_js.const import ( + CONF_DATA_COLLECTION_OPTED_IN, + DOMAIN, +) from homeassistant.helpers import device_registry as dr @@ -552,3 +556,72 @@ async def test_get_log_config(hass, client, integration, hass_ws_client): assert log_config["log_to_file"] is False assert log_config["filename"] == "/test.txt" assert log_config["force_console"] is False + + +async def test_data_collection(hass, client, integration, hass_ws_client): + """Test that the data collection WS API commands work.""" + entry = integration + ws_client = await hass_ws_client(hass) + + client.async_send_command.return_value = {"statisticsEnabled": False} + await ws_client.send_json( + { + ID: 1, + TYPE: "zwave_js/data_collection_status", + ENTRY_ID: entry.entry_id, + } + ) + msg = await ws_client.receive_json() + result = msg["result"] + assert result == {"opted_in": None, "enabled": False} + + assert len(client.async_send_command.call_args_list) == 1 + assert client.async_send_command.call_args[0][0] == { + "command": "driver.is_statistics_enabled" + } + + assert CONF_DATA_COLLECTION_OPTED_IN not in entry.data + + client.async_send_command.reset_mock() + + client.async_send_command.return_value = {} + await ws_client.send_json( + { + ID: 2, + TYPE: "zwave_js/update_data_collection_preference", + ENTRY_ID: entry.entry_id, + OPTED_IN: True, + } + ) + msg = await ws_client.receive_json() + result = msg["result"] + assert result is None + + assert len(client.async_send_command.call_args_list) == 1 + args = client.async_send_command.call_args[0][0] + assert args["command"] == "driver.enable_statistics" + assert args["applicationName"] == "Home Assistant" + assert entry.data[CONF_DATA_COLLECTION_OPTED_IN] + + client.async_send_command.reset_mock() + + client.async_send_command.return_value = {} + await ws_client.send_json( + { + ID: 3, + TYPE: "zwave_js/update_data_collection_preference", + ENTRY_ID: entry.entry_id, + OPTED_IN: False, + } + ) + msg = await ws_client.receive_json() + result = msg["result"] + assert result is None + + assert len(client.async_send_command.call_args_list) == 1 + assert client.async_send_command.call_args[0][0] == { + "command": "driver.disable_statistics" + } + assert not entry.data[CONF_DATA_COLLECTION_OPTED_IN] + + client.async_send_command.reset_mock() diff --git a/tests/components/zwave_js/test_init.py b/tests/components/zwave_js/test_init.py index 32fcdbcc84a..f9784e0f9b8 100644 --- a/tests/components/zwave_js/test_init.py +++ b/tests/components/zwave_js/test_init.py @@ -66,6 +66,54 @@ async def test_initialized_timeout(hass, client, connect_timeout): assert entry.state == ENTRY_STATE_SETUP_RETRY +async def test_enabled_statistics(hass, client): + """Test that we enabled statistics if the entry is opted in.""" + entry = MockConfigEntry( + domain="zwave_js", + data={"url": "ws://test.org", "data_collection_opted_in": True}, + ) + entry.add_to_hass(hass) + + with patch( + "zwave_js_server.model.driver.Driver.async_enable_statistics" + ) as mock_cmd: + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + assert mock_cmd.called + + +async def test_disabled_statistics(hass, client): + """Test that we diisabled statistics if the entry is opted out.""" + entry = MockConfigEntry( + domain="zwave_js", + data={"url": "ws://test.org", "data_collection_opted_in": False}, + ) + entry.add_to_hass(hass) + + with patch( + "zwave_js_server.model.driver.Driver.async_disable_statistics" + ) as mock_cmd: + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + assert mock_cmd.called + + +async def test_noop_statistics(hass, client): + """Test that we don't make any statistics calls if user hasn't provided preference.""" + entry = MockConfigEntry(domain="zwave_js", data={"url": "ws://test.org"}) + entry.add_to_hass(hass) + + with patch( + "zwave_js_server.model.driver.Driver.async_enable_statistics" + ) as mock_cmd1, patch( + "zwave_js_server.model.driver.Driver.async_disable_statistics" + ) as mock_cmd2: + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + assert not mock_cmd1.called + assert not mock_cmd2.called + + @pytest.mark.parametrize("error", [BaseZwaveJSServerError("Boom"), Exception("Boom")]) async def test_listen_failure(hass, client, error): """Test we handle errors during client listen."""