From 3663e0af4103524c0723d940f6b2efa020dd6639 Mon Sep 17 00:00:00 2001
From: Dave T <17680170+davet2001@users.noreply.github.com>
Date: Wed, 22 Dec 2021 08:16:37 +0000
Subject: [PATCH] Use enums for ozw tests (#62547)

---
 tests/components/ozw/test_binary_sensor.py |  4 ++--
 tests/components/ozw/test_sensor.py        | 17 ++++++-----------
 2 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/tests/components/ozw/test_binary_sensor.py b/tests/components/ozw/test_binary_sensor.py
index 6053cc43457..d0852a5caf0 100644
--- a/tests/components/ozw/test_binary_sensor.py
+++ b/tests/components/ozw/test_binary_sensor.py
@@ -1,7 +1,7 @@
 """Test Z-Wave Sensors."""
 from homeassistant.components.binary_sensor import (
-    DEVICE_CLASS_MOTION,
     DOMAIN as BINARY_SENSOR_DOMAIN,
+    BinarySensorDeviceClass,
 )
 from homeassistant.components.ozw.const import DOMAIN
 from homeassistant.const import ATTR_DEVICE_CLASS
@@ -35,7 +35,7 @@ async def test_binary_sensor(hass, generic_data, binary_sensor_msg):
     state = hass.states.get("binary_sensor.trisensor_home_security_motion_detected")
     assert state
     assert state.state == "off"
-    assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_MOTION
+    assert state.attributes[ATTR_DEVICE_CLASS] == BinarySensorDeviceClass.MOTION
 
     # Test incoming state change
     receive_msg(binary_sensor_msg)
diff --git a/tests/components/ozw/test_sensor.py b/tests/components/ozw/test_sensor.py
index 01bdb7b51a2..2eddb9722e5 100644
--- a/tests/components/ozw/test_sensor.py
+++ b/tests/components/ozw/test_sensor.py
@@ -1,11 +1,6 @@
 """Test Z-Wave Sensors."""
 from homeassistant.components.ozw.const import DOMAIN
-from homeassistant.components.sensor import (
-    DEVICE_CLASS_HUMIDITY,
-    DEVICE_CLASS_POWER,
-    DEVICE_CLASS_PRESSURE,
-    DOMAIN as SENSOR_DOMAIN,
-)
+from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN, SensorDeviceClass
 from homeassistant.const import ATTR_DEVICE_CLASS
 from homeassistant.helpers import entity_registry as er
 
@@ -24,15 +19,15 @@ async def test_sensor(hass, generic_data):
 
     # Test device classes
     state = hass.states.get("sensor.trisensor_relative_humidity")
-    assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_HUMIDITY
+    assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.HUMIDITY
     state = hass.states.get("sensor.trisensor_pressure")
-    assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_PRESSURE
+    assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.PRESSURE
     state = hass.states.get("sensor.trisensor_fake_power")
-    assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_POWER
+    assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.POWER
     state = hass.states.get("sensor.trisensor_fake_energy")
-    assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_POWER
+    assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.POWER
     state = hass.states.get("sensor.trisensor_fake_electric")
-    assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_POWER
+    assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.POWER
 
     # Test ZWaveListSensor disabled by default
     registry = er.async_get(hass)