From 3ebc5d45a8af9cb2e9d7d73d0bdd477443a6188c Mon Sep 17 00:00:00 2001
From: Petro31 <35082313+Petro31@users.noreply.github.com>
Date: Wed, 13 Jan 2021 15:37:54 -0500
Subject: [PATCH] Add energy and power sensor tests & fix device_class (#45122)

---
 homeassistant/components/zwave_js/sensor.py |  6 ++---
 tests/components/zwave_js/common.py         |  2 ++
 tests/components/zwave_js/test_sensor.py    | 28 +++++++++++++++++++--
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/zwave_js/sensor.py b/homeassistant/components/zwave_js/sensor.py
index ef1a68bb7a7..89c34844a1e 100644
--- a/homeassistant/components/zwave_js/sensor.py
+++ b/homeassistant/components/zwave_js/sensor.py
@@ -65,11 +65,9 @@ class ZwaveSensorBase(ZWaveBaseEntity):
         if self.info.primary_value.command_class == CommandClass.BATTERY:
             return DEVICE_CLASS_BATTERY
         if self.info.primary_value.command_class == CommandClass.METER:
+            if self.info.primary_value.property_key_name == "kWh_Consumed":
+                return DEVICE_CLASS_ENERGY
             return DEVICE_CLASS_POWER
-        if self.info.primary_value.property_key_name == "W_Consumed":
-            return DEVICE_CLASS_POWER
-        if self.info.primary_value.property_key_name == "kWh_Consumed":
-            return DEVICE_CLASS_ENERGY
         if self.info.primary_value.property_ == "Air temperature":
             return DEVICE_CLASS_TEMPERATURE
         return None
diff --git a/tests/components/zwave_js/common.py b/tests/components/zwave_js/common.py
index c78c6d544b6..73d70d31669 100644
--- a/tests/components/zwave_js/common.py
+++ b/tests/components/zwave_js/common.py
@@ -1,3 +1,5 @@
 """Provide common test tools for Z-Wave JS."""
 AIR_TEMPERATURE_SENSOR = "sensor.multisensor_6_air_temperature"
+ENERGY_SENSOR = "sensor.smart_plug_with_two_usb_ports_value_electric_consumed_2"
+POWER_SENSOR = "sensor.smart_plug_with_two_usb_ports_value_electric_consumed"
 SWITCH_ENTITY = "switch.smart_plug_with_two_usb_ports_current_value"
diff --git a/tests/components/zwave_js/test_sensor.py b/tests/components/zwave_js/test_sensor.py
index 75ce016fb04..284d2e1a84f 100644
--- a/tests/components/zwave_js/test_sensor.py
+++ b/tests/components/zwave_js/test_sensor.py
@@ -1,7 +1,14 @@
 """Test the Z-Wave JS sensor platform."""
-from homeassistant.const import DEVICE_CLASS_TEMPERATURE, TEMP_CELSIUS
+from homeassistant.const import (
+    DEVICE_CLASS_ENERGY,
+    DEVICE_CLASS_POWER,
+    DEVICE_CLASS_TEMPERATURE,
+    ENERGY_KILO_WATT_HOUR,
+    POWER_WATT,
+    TEMP_CELSIUS,
+)
 
-from .common import AIR_TEMPERATURE_SENSOR
+from .common import AIR_TEMPERATURE_SENSOR, ENERGY_SENSOR, POWER_SENSOR
 
 
 async def test_numeric_sensor(hass, multisensor_6, integration):
@@ -12,3 +19,20 @@ async def test_numeric_sensor(hass, multisensor_6, integration):
     assert state.state == "9.0"
     assert state.attributes["unit_of_measurement"] == TEMP_CELSIUS
     assert state.attributes["device_class"] == DEVICE_CLASS_TEMPERATURE
+
+
+async def test_energy_sensors(hass, hank_binary_switch, integration):
+    """Test power and energy sensors."""
+    state = hass.states.get(POWER_SENSOR)
+
+    assert state
+    assert state.state == "0.0"
+    assert state.attributes["unit_of_measurement"] == POWER_WATT
+    assert state.attributes["device_class"] == DEVICE_CLASS_POWER
+
+    state = hass.states.get(ENERGY_SENSOR)
+
+    assert state
+    assert state.state == "0.16"
+    assert state.attributes["unit_of_measurement"] == ENERGY_KILO_WATT_HOUR
+    assert state.attributes["device_class"] == DEVICE_CLASS_ENERGY