Get zwave_js statistics data from model (#120281)

* Get zwave_js statistics data from model

* Add migration logic

* Update comment

* revert change to forward entry

homeassistant/components/zwave_js/__init__.py

@@ -353,7 +353,7 @@ class ControllerEvents:
         self.discovered_value_ids: dict[str, set[str]] = defaultdict(set)
         self.driver_events = driver_events
         self.dev_reg = driver_events.dev_reg
-        self.registered_unique_ids: dict[str, dict[str, set[str]]] = defaultdict(
+        self.registered_unique_ids: dict[str, dict[Platform, set[str]]] = defaultdict(
             lambda: defaultdict(set)
         )
         self.node_events = NodeEvents(hass, self)

homeassistant/components/zwave_js/migrate.py

@@ -6,20 +6,16 @@ from dataclasses import dataclass
 import logging
 
 from zwave_js_server.model.driver import Driver
+from zwave_js_server.model.node import Node
 from zwave_js_server.model.value import Value as ZwaveValue
 
-from homeassistant.const import STATE_UNAVAILABLE
+from homeassistant.const import STATE_UNAVAILABLE, Platform
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.device_registry import DeviceEntry
+from homeassistant.helpers import device_registry as dr, entity_registry as er
-from homeassistant.helpers.entity_registry import (
-    EntityRegistry,
-    RegistryEntry,
-    async_entries_for_device,
-)
 
 from .const import DOMAIN
 from .discovery import ZwaveDiscoveryInfo
-from .helpers import get_unique_id
+from .helpers import get_unique_id, get_valueless_base_unique_id
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -62,10 +58,10 @@ class ValueID:
 @callback
 def async_migrate_old_entity(
     hass: HomeAssistant,
-    ent_reg: EntityRegistry,
+    ent_reg: er.EntityRegistry,
     registered_unique_ids: set[str],
-    platform: str,
+    platform: Platform,
-    device: DeviceEntry,
+    device: dr.DeviceEntry,
     unique_id: str,
 ) -> None:
     """Migrate existing entity if current one can't be found and an old one exists."""
@@ -77,8 +73,8 @@ def async_migrate_old_entity(
 
     # Look for existing entities in the registry that could be the same value but on
     # a different endpoint
-    existing_entity_entries: list[RegistryEntry] = []
+    existing_entity_entries: list[er.RegistryEntry] = []
-    for entry in async_entries_for_device(ent_reg, device.id):
+    for entry in er.async_entries_for_device(ent_reg, device.id):
         # If entity is not in the domain for this discovery info or entity has already
         # been processed, skip it
         if entry.domain != platform or entry.unique_id in registered_unique_ids:
@@ -109,35 +105,40 @@ def async_migrate_old_entity(
 
 @callback
 def async_migrate_unique_id(
-    ent_reg: EntityRegistry, platform: str, old_unique_id: str, new_unique_id: str
+    ent_reg: er.EntityRegistry,
+    platform: Platform,
+    old_unique_id: str,
+    new_unique_id: str,
 ) -> None:
     """Check if entity with old unique ID exists, and if so migrate it to new ID."""
-    if entity_id := ent_reg.async_get_entity_id(platform, DOMAIN, old_unique_id):
+    if not (entity_id := ent_reg.async_get_entity_id(platform, DOMAIN, old_unique_id)):
+        return
+
+    _LOGGER.debug(
+        "Migrating entity %s from old unique ID '%s' to new unique ID '%s'",
+        entity_id,
+        old_unique_id,
+        new_unique_id,
+    )
+    try:
+        ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id)
+    except ValueError:
         _LOGGER.debug(
-            "Migrating entity %s from old unique ID '%s' to new unique ID '%s'",
+            (
+                "Entity %s can't be migrated because the unique ID is taken; "
+                "Cleaning it up since it is likely no longer valid"
+            ),
             entity_id,
-            old_unique_id,
-            new_unique_id,
         )
-        try:
+        ent_reg.async_remove(entity_id)
-            ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id)
-        except ValueError:
-            _LOGGER.debug(
-                (
-                    "Entity %s can't be migrated because the unique ID is taken; "
-                    "Cleaning it up since it is likely no longer valid"
-                ),
-                entity_id,
-            )
-            ent_reg.async_remove(entity_id)
 
 
 @callback
 def async_migrate_discovered_value(
     hass: HomeAssistant,
-    ent_reg: EntityRegistry,
+    ent_reg: er.EntityRegistry,
     registered_unique_ids: set[str],
-    device: DeviceEntry,
+    device: dr.DeviceEntry,
     driver: Driver,
     disc_info: ZwaveDiscoveryInfo,
 ) -> None:
@@ -160,7 +161,7 @@ def async_migrate_discovered_value(
     ]
 
     if (
-        disc_info.platform == "binary_sensor"
+        disc_info.platform == Platform.BINARY_SENSOR
         and disc_info.platform_hint == "notification"
     ):
         for state_key in disc_info.primary_value.metadata.states:
@@ -211,6 +212,24 @@ def async_migrate_discovered_value(
     registered_unique_ids.add(new_unique_id)
 
 
+@callback
+def async_migrate_statistics_sensors(
+    hass: HomeAssistant, driver: Driver, node: Node, key_map: dict[str, str]
+) -> None:
+    """Migrate statistics sensors to new unique IDs.
+
+    - Migrate camel case keys in unique IDs to snake keys.
+    """
+    ent_reg = er.async_get(hass)
+    base_unique_id = f"{get_valueless_base_unique_id(driver, node)}.statistics"
+    for new_key, old_key in key_map.items():
+        if new_key == old_key:
+            continue
+        old_unique_id = f"{base_unique_id}_{old_key}"
+        new_unique_id = f"{base_unique_id}_{new_key}"
+        async_migrate_unique_id(ent_reg, Platform.SENSOR, old_unique_id, new_unique_id)
+
+
 @callback
 def get_old_value_ids(value: ZwaveValue) -> list[str]:
     """Get old value IDs so we can migrate entity unique ID."""

homeassistant/components/zwave_js/sensor.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
-from datetime import datetime
 from typing import Any
 
 import voluptuous as vol
@@ -16,10 +15,10 @@ from zwave_js_server.const.command_class.meter import (
 )
 from zwave_js_server.exceptions import BaseZwaveJSServerError
 from zwave_js_server.model.controller import Controller
-from zwave_js_server.model.controller.statistics import ControllerStatisticsDataType
+from zwave_js_server.model.controller.statistics import ControllerStatistics
 from zwave_js_server.model.driver import Driver
 from zwave_js_server.model.node import Node as ZwaveNode
-from zwave_js_server.model.node.statistics import NodeStatisticsDataType
+from zwave_js_server.model.node.statistics import NodeStatistics
 from zwave_js_server.util.command_class.meter import get_meter_type
 
 from homeassistant.components.sensor import (
@@ -90,6 +89,7 @@ from .discovery_data_template import (
 )
 from .entity import ZWaveBaseEntity
 from .helpers import get_device_info, get_valueless_base_unique_id
+from .migrate import async_migrate_statistics_sensors
 
 PARALLEL_UPDATES = 0
 
@@ -328,152 +328,172 @@ ENTITY_DESCRIPTION_KEY_MAP = {
 }
 
 
-def convert_dict_of_dicts(
+def convert_nested_attr(
-    statistics: ControllerStatisticsDataType | NodeStatisticsDataType, key: str
+    statistics: ControllerStatistics | NodeStatistics, key: str
 ) -> Any:
-    """Convert a dictionary of dictionaries to a value."""
+    """Convert a string that represents a nested attr to a value."""
-    keys = key.split(".")
+    data = statistics
-    return statistics.get(keys[0], {}).get(keys[1], {}).get(keys[2])  # type: ignore[attr-defined]
+    for _key in key.split("."):
+        if data is None:
+            return None  # type: ignore[unreachable]
+        data = getattr(data, _key)
+    return data
 
 
 @dataclass(frozen=True, kw_only=True)
 class ZWaveJSStatisticsSensorEntityDescription(SensorEntityDescription):
     """Class to represent a Z-Wave JS statistics sensor entity description."""
 
-    convert: Callable[
+    convert: Callable[[ControllerStatistics | NodeStatistics, str], Any] = getattr
-        [ControllerStatisticsDataType | NodeStatisticsDataType, str], Any
-    ] = lambda statistics, key: statistics.get(key)
     entity_registry_enabled_default: bool = False
 
 
 # Controller statistics descriptions
 ENTITY_DESCRIPTION_CONTROLLER_STATISTICS_LIST = [
     ZWaveJSStatisticsSensorEntityDescription(
-        key="messagesTX",
+        key="messages_tx",
         translation_key="successful_messages",
         translation_placeholders={"direction": "TX"},
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="messagesRX",
+        key="messages_rx",
         translation_key="successful_messages",
         translation_placeholders={"direction": "RX"},
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="messagesDroppedTX",
+        key="messages_dropped_tx",
         translation_key="messages_dropped",
         translation_placeholders={"direction": "TX"},
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="messagesDroppedRX",
+        key="messages_dropped_rx",
         translation_key="messages_dropped",
         translation_placeholders={"direction": "RX"},
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="NAK", translation_key="nak", state_class=SensorStateClass.TOTAL
+        key="nak", translation_key="nak", state_class=SensorStateClass.TOTAL
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="CAN", translation_key="can", state_class=SensorStateClass.TOTAL
+        key="can", translation_key="can", state_class=SensorStateClass.TOTAL
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="timeoutACK",
+        key="timeout_ack",
         translation_key="timeout_ack",
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="timeoutResponse",
+        key="timeout_response",
         translation_key="timeout_response",
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="timeoutCallback",
+        key="timeout_callback",
         translation_key="timeout_callback",
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="backgroundRSSI.channel0.average",
+        key="background_rssi.channel_0.average",
         translation_key="average_background_rssi",
         translation_placeholders={"channel": "0"},
         native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
         device_class=SensorDeviceClass.SIGNAL_STRENGTH,
-        convert=convert_dict_of_dicts,
+        convert=convert_nested_attr,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="backgroundRSSI.channel0.current",
+        key="background_rssi.channel_0.current",
         translation_key="current_background_rssi",
         translation_placeholders={"channel": "0"},
         native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
         device_class=SensorDeviceClass.SIGNAL_STRENGTH,
         state_class=SensorStateClass.MEASUREMENT,
-        convert=convert_dict_of_dicts,
+        convert=convert_nested_attr,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="backgroundRSSI.channel1.average",
+        key="background_rssi.channel_1.average",
         translation_key="average_background_rssi",
         translation_placeholders={"channel": "1"},
         native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
         device_class=SensorDeviceClass.SIGNAL_STRENGTH,
-        convert=convert_dict_of_dicts,
+        convert=convert_nested_attr,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="backgroundRSSI.channel1.current",
+        key="background_rssi.channel_1.current",
         translation_key="current_background_rssi",
         translation_placeholders={"channel": "1"},
         native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
         device_class=SensorDeviceClass.SIGNAL_STRENGTH,
         state_class=SensorStateClass.MEASUREMENT,
-        convert=convert_dict_of_dicts,
+        convert=convert_nested_attr,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="backgroundRSSI.channel2.average",
+        key="background_rssi.channel_2.average",
         translation_key="average_background_rssi",
         translation_placeholders={"channel": "2"},
         native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
         device_class=SensorDeviceClass.SIGNAL_STRENGTH,
-        convert=convert_dict_of_dicts,
+        convert=convert_nested_attr,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="backgroundRSSI.channel2.current",
+        key="background_rssi.channel_2.current",
         translation_key="current_background_rssi",
         translation_placeholders={"channel": "2"},
         native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
         device_class=SensorDeviceClass.SIGNAL_STRENGTH,
         state_class=SensorStateClass.MEASUREMENT,
-        convert=convert_dict_of_dicts,
+        convert=convert_nested_attr,
     ),
 ]
 
+CONTROLLER_STATISTICS_KEY_MAP: dict[str, str] = {
+    "messages_tx": "messagesTX",
+    "messages_rx": "messagesRX",
+    "messages_dropped_tx": "messagesDroppedTX",
+    "messages_dropped_rx": "messagesDroppedRX",
+    "nak": "NAK",
+    "can": "CAN",
+    "timeout_ack": "timeoutAck",
+    "timeout_response": "timeoutResponse",
+    "timeout_callback": "timeoutCallback",
+    "background_rssi.channel_0.average": "backgroundRSSI.channel0.average",
+    "background_rssi.channel_0.current": "backgroundRSSI.channel0.current",
+    "background_rssi.channel_1.average": "backgroundRSSI.channel1.average",
+    "background_rssi.channel_1.current": "backgroundRSSI.channel1.current",
+    "background_rssi.channel_2.average": "backgroundRSSI.channel2.average",
+    "background_rssi.channel_2.current": "backgroundRSSI.channel2.current",
+}
+
 # Node statistics descriptions
 ENTITY_DESCRIPTION_NODE_STATISTICS_LIST = [
     ZWaveJSStatisticsSensorEntityDescription(
-        key="commandsRX",
+        key="commands_rx",
         translation_key="successful_commands",
         translation_placeholders={"direction": "RX"},
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="commandsTX",
+        key="commands_tx",
         translation_key="successful_commands",
         translation_placeholders={"direction": "TX"},
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="commandsDroppedRX",
+        key="commands_dropped_rx",
         translation_key="commands_dropped",
         translation_placeholders={"direction": "RX"},
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="commandsDroppedTX",
+        key="commands_dropped_tx",
         translation_key="commands_dropped",
         translation_placeholders={"direction": "TX"},
         state_class=SensorStateClass.TOTAL,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="timeoutResponse",
+        key="timeout_response",
         translation_key="timeout_response",
         state_class=SensorStateClass.TOTAL,
     ),
@@ -492,20 +512,24 @@ ENTITY_DESCRIPTION_NODE_STATISTICS_LIST = [
         state_class=SensorStateClass.MEASUREMENT,
     ),
     ZWaveJSStatisticsSensorEntityDescription(
-        key="lastSeen",
+        key="last_seen",
         translation_key="last_seen",
         device_class=SensorDeviceClass.TIMESTAMP,
-        convert=(
-            lambda statistics, key: (
-                datetime.fromisoformat(dt)  # type: ignore[arg-type]
-                if (dt := statistics.get(key))
-                else None
-            )
-        ),
         entity_registry_enabled_default=True,
     ),
 ]
 
+NODE_STATISTICS_KEY_MAP: dict[str, str] = {
+    "commands_rx": "commandsRX",
+    "commands_tx": "commandsTX",
+    "commands_dropped_rx": "commandsDroppedRX",
+    "commands_dropped_tx": "commandsDroppedTX",
+    "timeout_response": "timeoutResponse",
+    "rtt": "rtt",
+    "rssi": "rssi",
+    "last_seen": "lastSeen",
+}
+
 
 def get_entity_description(
     data: NumericSensorDataTemplateData,
@@ -588,6 +612,14 @@ async def async_setup_entry(
     @callback
     def async_add_statistics_sensors(node: ZwaveNode) -> None:
         """Add statistics sensors."""
+        async_migrate_statistics_sensors(
+            hass,
+            driver,
+            node,
+            CONTROLLER_STATISTICS_KEY_MAP
+            if driver.controller.own_node == node
+            else NODE_STATISTICS_KEY_MAP,
+        )
         async_add_entities(
             [
                 ZWaveStatisticsSensor(
@@ -1001,7 +1033,7 @@ class ZWaveStatisticsSensor(SensorEntity):
     def statistics_updated(self, event_data: dict) -> None:
         """Call when statistics updated event is received."""
         self._attr_native_value = self.entity_description.convert(
-            event_data["statistics"], self.entity_description.key
+            event_data["statistics_updated"], self.entity_description.key
         )
         self.async_write_ha_state()
 
@@ -1027,5 +1059,5 @@ class ZWaveStatisticsSensor(SensorEntity):
 
         # Set initial state
         self._attr_native_value = self.entity_description.convert(
-            self.statistics_src.statistics.data, self.entity_description.key
+            self.statistics_src.statistics, self.entity_description.key
         )

tests/components/zwave_js/test_sensor.py

@@ -23,6 +23,10 @@ from homeassistant.components.zwave_js.const import (
     SERVICE_RESET_METER,
 )
 from homeassistant.components.zwave_js.helpers import get_valueless_base_unique_id
+from homeassistant.components.zwave_js.sensor import (
+    CONTROLLER_STATISTICS_KEY_MAP,
+    NODE_STATISTICS_KEY_MAP,
+)
 from homeassistant.const import (
     ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
@@ -55,6 +59,8 @@ from .common import (
     VOLTAGE_SENSOR,
 )
 
+from tests.common import MockConfigEntry
+
 
 async def test_numeric_sensor(
     hass: HomeAssistant,
@@ -756,6 +762,54 @@ NODE_STATISTICS_SUFFIXES_UNKNOWN = {
 }
 
 
+async def test_statistics_sensors_migration(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    zp3111_state,
+    client,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test statistics migration sensor."""
+    node = Node(client, copy.deepcopy(zp3111_state))
+    client.driver.controller.nodes[node.node_id] = node
+
+    entry = MockConfigEntry(domain="zwave_js", data={"url": "ws://test.org"})
+    entry.add_to_hass(hass)
+
+    controller_base_unique_id = f"{client.driver.controller.home_id}.1.statistics"
+    node_base_unique_id = f"{client.driver.controller.home_id}.22.statistics"
+
+    # Create entity registry records for the old statistics keys
+    for base_unique_id, key_map in (
+        (controller_base_unique_id, CONTROLLER_STATISTICS_KEY_MAP),
+        (node_base_unique_id, NODE_STATISTICS_KEY_MAP),
+    ):
+        # old key
+        for key in key_map.values():
+            entity_registry.async_get_or_create(
+                "sensor", DOMAIN, f"{base_unique_id}_{key}"
+            )
+
+    # Set up integration
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Validate that entity unique ID's have changed
+    for base_unique_id, key_map in (
+        (controller_base_unique_id, CONTROLLER_STATISTICS_KEY_MAP),
+        (node_base_unique_id, NODE_STATISTICS_KEY_MAP),
+    ):
+        for new_key, old_key in key_map.items():
+            # If the key has changed, the old entity should not exist
+            if new_key != old_key:
+                assert not entity_registry.async_get_entity_id(
+                    "sensor", DOMAIN, f"{base_unique_id}_{old_key}"
+                )
+            assert entity_registry.async_get_entity_id(
+                "sensor", DOMAIN, f"{base_unique_id}_{new_key}"
+            )
+
+
 async def test_statistics_sensors_no_last_seen(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,