From d62a78ae6196226f24a5bb5a1240c2f1bf787be8 Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Mon, 30 Aug 2021 11:48:36 -0400
Subject: [PATCH] Don't set zwave_js sensor device class to energy when unit is
 wrong (#55434)

---
 homeassistant/components/zwave_js/const.py       |  2 ++
 .../zwave_js/discovery_data_template.py          | 16 ++++++++++++++++
 homeassistant/components/zwave_js/sensor.py      | 12 ++++++++++++
 3 files changed, 30 insertions(+)

diff --git a/homeassistant/components/zwave_js/const.py b/homeassistant/components/zwave_js/const.py
index 8e545975faa..e4486a681e1 100644
--- a/homeassistant/components/zwave_js/const.py
+++ b/homeassistant/components/zwave_js/const.py
@@ -96,3 +96,5 @@ ENTITY_DESC_KEY_SIGNAL_STRENGTH = "signal_strength"
 ENTITY_DESC_KEY_TEMPERATURE = "temperature"
 ENTITY_DESC_KEY_TARGET_TEMPERATURE = "target_temperature"
 ENTITY_DESC_KEY_TIMESTAMP = "timestamp"
+ENTITY_DESC_KEY_MEASUREMENT = "measurement"
+ENTITY_DESC_KEY_TOTAL_INCREASING = "total_increasing"
diff --git a/homeassistant/components/zwave_js/discovery_data_template.py b/homeassistant/components/zwave_js/discovery_data_template.py
index 3ef74a7e17d..dd338de63eb 100644
--- a/homeassistant/components/zwave_js/discovery_data_template.py
+++ b/homeassistant/components/zwave_js/discovery_data_template.py
@@ -24,6 +24,7 @@ from zwave_js_server.const import (
     VOLTAGE_METER_TYPES,
     VOLTAGE_SENSORS,
     CommandClass,
+    ElectricScale,
     MeterScaleType,
     MultilevelSensorType,
 )
@@ -43,6 +44,7 @@ from .const import (
     ENTITY_DESC_KEY_ENERGY_TOTAL_INCREASING,
     ENTITY_DESC_KEY_HUMIDITY,
     ENTITY_DESC_KEY_ILLUMINANCE,
+    ENTITY_DESC_KEY_MEASUREMENT,
     ENTITY_DESC_KEY_POWER,
     ENTITY_DESC_KEY_POWER_FACTOR,
     ENTITY_DESC_KEY_PRESSURE,
@@ -50,6 +52,7 @@ from .const import (
     ENTITY_DESC_KEY_TARGET_TEMPERATURE,
     ENTITY_DESC_KEY_TEMPERATURE,
     ENTITY_DESC_KEY_TIMESTAMP,
+    ENTITY_DESC_KEY_TOTAL_INCREASING,
     ENTITY_DESC_KEY_VOLTAGE,
 )
 
@@ -187,6 +190,19 @@ class NumericSensorDataTemplate(BaseDiscoverySchemaDataTemplate):
 
         if value.command_class == CommandClass.METER:
             scale_type = get_meter_scale_type(value)
+            # We do this because even though these are energy scales, they don't meet
+            # the unit requirements for the energy device class.
+            if scale_type in (
+                ElectricScale.PULSE,
+                ElectricScale.KILOVOLT_AMPERE_HOUR,
+                ElectricScale.KILOVOLT_AMPERE_REACTIVE_HOUR,
+            ):
+                return ENTITY_DESC_KEY_TOTAL_INCREASING
+            # We do this because even though these are power scales, they don't meet
+            # the unit requirements for the energy power class.
+            if scale_type == ElectricScale.KILOVOLT_AMPERE_REACTIVE:
+                return ENTITY_DESC_KEY_MEASUREMENT
+
             for key, scale_type_set in METER_DEVICE_CLASS_MAP.items():
                 if scale_type in scale_type_set:
                     return key
diff --git a/homeassistant/components/zwave_js/sensor.py b/homeassistant/components/zwave_js/sensor.py
index 40159b383a6..c71a1d87653 100644
--- a/homeassistant/components/zwave_js/sensor.py
+++ b/homeassistant/components/zwave_js/sensor.py
@@ -63,6 +63,7 @@ from .const import (
     ENTITY_DESC_KEY_ENERGY_TOTAL_INCREASING,
     ENTITY_DESC_KEY_HUMIDITY,
     ENTITY_DESC_KEY_ILLUMINANCE,
+    ENTITY_DESC_KEY_MEASUREMENT,
     ENTITY_DESC_KEY_POWER,
     ENTITY_DESC_KEY_POWER_FACTOR,
     ENTITY_DESC_KEY_PRESSURE,
@@ -70,6 +71,7 @@ from .const import (
     ENTITY_DESC_KEY_TARGET_TEMPERATURE,
     ENTITY_DESC_KEY_TEMPERATURE,
     ENTITY_DESC_KEY_TIMESTAMP,
+    ENTITY_DESC_KEY_TOTAL_INCREASING,
     ENTITY_DESC_KEY_VOLTAGE,
     SERVICE_RESET_METER,
 )
@@ -168,6 +170,16 @@ ENTITY_DESCRIPTION_KEY_MAP: dict[str, ZwaveSensorEntityDescription] = {
         device_class=DEVICE_CLASS_TEMPERATURE,
         state_class=None,
     ),
+    ENTITY_DESC_KEY_MEASUREMENT: ZwaveSensorEntityDescription(
+        ENTITY_DESC_KEY_MEASUREMENT,
+        device_class=None,
+        state_class=STATE_CLASS_MEASUREMENT,
+    ),
+    ENTITY_DESC_KEY_TOTAL_INCREASING: ZwaveSensorEntityDescription(
+        ENTITY_DESC_KEY_TOTAL_INCREASING,
+        device_class=None,
+        state_class=STATE_CLASS_TOTAL_INCREASING,
+    ),
 }