Fix KeyError from zwave_js diagnostics (#74579)

homeassistant/components/zwave_js/diagnostics.py

@@ -94,20 +94,23 @@ def get_device_entities(
         # If the value ID returns as None, we don't need to include this entity
         if (value_id := get_value_id_from_unique_id(entry.unique_id)) is None:
             continue
-        state_key = get_state_key_from_unique_id(entry.unique_id)
 
-        zwave_value = node.values[value_id]
-        primary_value_data = {
-            "command_class": zwave_value.command_class,
-            "command_class_name": zwave_value.command_class_name,
-            "endpoint": zwave_value.endpoint,
-            "property": zwave_value.property_,
-            "property_name": zwave_value.property_name,
-            "property_key": zwave_value.property_key,
-            "property_key_name": zwave_value.property_key_name,
-        }
-        if state_key is not None:
-            primary_value_data["state_key"] = state_key
+        primary_value_data = None
+        if (zwave_value := node.values.get(value_id)) is not None:
+            primary_value_data = {
+                "command_class": zwave_value.command_class,
+                "command_class_name": zwave_value.command_class_name,
+                "endpoint": zwave_value.endpoint,
+                "property": zwave_value.property_,
+                "property_name": zwave_value.property_name,
+                "property_key": zwave_value.property_key,
+                "property_key_name": zwave_value.property_key_name,
+            }
+
+            state_key = get_state_key_from_unique_id(entry.unique_id)
+            if state_key is not None:
+                primary_value_data["state_key"] = state_key
+
         entity = {
             "domain": entry.domain,
             "entity_id": entry.entity_id,

tests/components/zwave_js/test_diagnostics.py

@@ -10,8 +10,12 @@ from homeassistant.components.zwave_js.diagnostics import (
     async_get_device_diagnostics,
 )
 from homeassistant.components.zwave_js.discovery import async_discover_node_values
-from homeassistant.components.zwave_js.helpers import get_device_id
-from homeassistant.helpers.device_registry import async_get
+from homeassistant.components.zwave_js.helpers import (
+    get_device_id,
+    get_value_id_from_unique_id,
+)
+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 .common import PROPERTY_ULTRAVIOLET
 
@@ -53,7 +57,7 @@ async def test_device_diagnostics(
     version_state,
 ):
     """Test the device level diagnostics data dump."""
-    dev_reg = async_get(hass)
+    dev_reg = async_get_dev_reg(hass)
     device = dev_reg.async_get_device({get_device_id(client.driver, multisensor_6)})
     assert device
 
@@ -106,7 +110,7 @@ async def test_device_diagnostics(
 
 async def test_device_diagnostics_error(hass, integration):
     """Test the device diagnostics raises exception when an invalid device is used."""
-    dev_reg = async_get(hass)
+    dev_reg = async_get_dev_reg(hass)
     device = dev_reg.async_get_or_create(
         config_entry_id=integration.entry_id, identifiers={("test", "test")}
     )
@@ -118,3 +122,71 @@ async def test_empty_zwave_value_matcher():
     """Test empty ZwaveValueMatcher is invalid."""
     with pytest.raises(ValueError):
         ZwaveValueMatcher()
+
+
+async def test_device_diagnostics_missing_primary_value(
+    hass,
+    client,
+    multisensor_6,
+    integration,
+    hass_client,
+):
+    """Test that the device diagnostics handles an entity with a missing primary value."""
+    dev_reg = async_get_dev_reg(hass)
+    device = dev_reg.async_get_device({get_device_id(client.driver, multisensor_6)})
+    assert device
+
+    entity_id = "sensor.multisensor_6_air_temperature"
+    ent_reg = async_get_ent_reg(hass)
+    entry = ent_reg.async_get(entity_id)
+
+    # check that the primary value for the entity exists in the diagnostics
+    diagnostics_data = await get_diagnostics_for_device(
+        hass, hass_client, integration, device
+    )
+
+    value = multisensor_6.values.get(get_value_id_from_unique_id(entry.unique_id))
+    assert value
+
+    air_entity = next(
+        x for x in diagnostics_data["entities"] if x["entity_id"] == entity_id
+    )
+
+    assert air_entity["primary_value"] == {
+        "command_class": value.command_class,
+        "command_class_name": value.command_class_name,
+        "endpoint": value.endpoint,
+        "property": value.property_,
+        "property_name": value.property_name,
+        "property_key": value.property_key,
+        "property_key_name": value.property_key_name,
+    }
+
+    # make the entity's primary value go missing
+    event = Event(
+        type="value removed",
+        data={
+            "source": "node",
+            "event": "value removed",
+            "nodeId": multisensor_6.node_id,
+            "args": {
+                "commandClassName": value.command_class_name,
+                "commandClass": value.command_class,
+                "endpoint": value.endpoint,
+                "property": value.property_,
+                "prevValue": 0,
+                "propertyName": value.property_name,
+            },
+        },
+    )
+    multisensor_6.receive_event(event)
+
+    diagnostics_data = await get_diagnostics_for_device(
+        hass, hass_client, integration, device
+    )
+
+    air_entity = next(
+        x for x in diagnostics_data["entities"] if x["entity_id"] == entity_id
+    )
+
+    assert air_entity["primary_value"] is None