Add support to enable/disable zwave_js data collection (#49440)

This commit is contained in:
Raman Gupta 2021-04-20 21:40:54 -04:00 committed by GitHub
parent a90d3a051f
commit 6e22251e1d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 287 additions and 28 deletions

View file

@ -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

View file

@ -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."""

View file

@ -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"

View file

@ -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

View file

@ -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()

View file

@ -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."""