From 5642350070a680bff22a189c50bb3281c2184fb9 Mon Sep 17 00:00:00 2001
From: Martin Hjelmare <marhje52@gmail.com>
Date: Mon, 25 Oct 2021 18:14:44 +0200
Subject: [PATCH] Add zwave_js sensor entity categories (#58416)

---
 homeassistant/components/zwave_js/sensor.py |  4 ++++
 tests/components/zwave_js/common.py         |  1 +
 tests/components/zwave_js/test_sensor.py    | 15 +++++++++++++++
 3 files changed, 20 insertions(+)

diff --git a/homeassistant/components/zwave_js/sensor.py b/homeassistant/components/zwave_js/sensor.py
index 7da7ba9ab9b..815158ee967 100644
--- a/homeassistant/components/zwave_js/sensor.py
+++ b/homeassistant/components/zwave_js/sensor.py
@@ -40,6 +40,7 @@ from homeassistant.const import (
     DEVICE_CLASS_TEMPERATURE,
     DEVICE_CLASS_TIMESTAMP,
     DEVICE_CLASS_VOLTAGE,
+    ENTITY_CATEGORY_DIAGNOSTIC,
     TEMP_CELSIUS,
     TEMP_FAHRENHEIT,
 )
@@ -100,6 +101,7 @@ ENTITY_DESCRIPTION_KEY_MAP: dict[str, ZwaveSensorEntityDescription] = {
     ENTITY_DESC_KEY_BATTERY: ZwaveSensorEntityDescription(
         ENTITY_DESC_KEY_BATTERY,
         device_class=DEVICE_CLASS_BATTERY,
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
         state_class=STATE_CLASS_MEASUREMENT,
     ),
     ENTITY_DESC_KEY_CURRENT: ZwaveSensorEntityDescription(
@@ -160,6 +162,8 @@ ENTITY_DESCRIPTION_KEY_MAP: dict[str, ZwaveSensorEntityDescription] = {
     ENTITY_DESC_KEY_SIGNAL_STRENGTH: ZwaveSensorEntityDescription(
         ENTITY_DESC_KEY_SIGNAL_STRENGTH,
         device_class=DEVICE_CLASS_SIGNAL_STRENGTH,
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        entity_registry_enabled_default=False,
         state_class=STATE_CLASS_MEASUREMENT,
     ),
     ENTITY_DESC_KEY_TEMPERATURE: ZwaveSensorEntityDescription(
diff --git a/tests/components/zwave_js/common.py b/tests/components/zwave_js/common.py
index e8e3151134c..c16ab00b2eb 100644
--- a/tests/components/zwave_js/common.py
+++ b/tests/components/zwave_js/common.py
@@ -1,5 +1,6 @@
 """Provide common test tools for Z-Wave JS."""
 AIR_TEMPERATURE_SENSOR = "sensor.multisensor_6_air_temperature"
+BATTERY_SENSOR = "sensor.multisensor_6_battery_level"
 HUMIDITY_SENSOR = "sensor.multisensor_6_humidity"
 POWER_SENSOR = "sensor.smart_plug_with_two_usb_ports_value_electric_consumed"
 ENERGY_SENSOR = "sensor.smart_plug_with_two_usb_ports_value_electric_consumed_2"
diff --git a/tests/components/zwave_js/test_sensor.py b/tests/components/zwave_js/test_sensor.py
index a18bead36c2..9e77d877b5c 100644
--- a/tests/components/zwave_js/test_sensor.py
+++ b/tests/components/zwave_js/test_sensor.py
@@ -21,6 +21,7 @@ from homeassistant.const import (
     ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
     ATTR_ICON,
+    DEVICE_CLASS_BATTERY,
     DEVICE_CLASS_CURRENT,
     DEVICE_CLASS_ENERGY,
     DEVICE_CLASS_HUMIDITY,
@@ -30,6 +31,7 @@ from homeassistant.const import (
     ELECTRIC_CURRENT_AMPERE,
     ELECTRIC_POTENTIAL_VOLT,
     ENERGY_KILO_WATT_HOUR,
+    ENTITY_CATEGORY_DIAGNOSTIC,
     POWER_WATT,
     STATE_UNAVAILABLE,
     TEMP_CELSIUS,
@@ -38,6 +40,7 @@ from homeassistant.helpers import entity_registry as er
 
 from .common import (
     AIR_TEMPERATURE_SENSOR,
+    BATTERY_SENSOR,
     CURRENT_SENSOR,
     ENERGY_SENSOR,
     HUMIDITY_SENSOR,
@@ -59,6 +62,18 @@ async def test_numeric_sensor(hass, multisensor_6, integration):
     assert state.attributes["unit_of_measurement"] == TEMP_CELSIUS
     assert state.attributes["device_class"] == DEVICE_CLASS_TEMPERATURE
 
+    state = hass.states.get(BATTERY_SENSOR)
+
+    assert state
+    assert state.state == "100.0"
+    assert state.attributes["unit_of_measurement"] == "%"
+    assert state.attributes["device_class"] == DEVICE_CLASS_BATTERY
+
+    ent_reg = er.async_get(hass)
+    entity_entry = ent_reg.async_get(BATTERY_SENSOR)
+    assert entity_entry
+    assert entity_entry.entity_category == ENTITY_CATEGORY_DIAGNOSTIC
+
     state = hass.states.get(HUMIDITY_SENSOR)
 
     assert state