From 0409665907e5db5add7345a7b30e07e4754ca95a Mon Sep 17 00:00:00 2001
From: Dave T <17680170+davet2001@users.noreply.github.com>
Date: Thu, 16 Dec 2021 20:22:26 +0000
Subject: [PATCH] Use DeviceClass Enums in ambee tests (#62108)

---
 tests/components/ambee/test_sensor.py | 28 +++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/tests/components/ambee/test_sensor.py b/tests/components/ambee/test_sensor.py
index 32bf5938456..a32398139d0 100644
--- a/tests/components/ambee/test_sensor.py
+++ b/tests/components/ambee/test_sensor.py
@@ -7,7 +7,8 @@ from homeassistant.components.ambee.const import DEVICE_CLASS_AMBEE_RISK, DOMAIN
 from homeassistant.components.sensor import (
     ATTR_STATE_CLASS,
     DOMAIN as SENSOR_DOMAIN,
-    STATE_CLASS_MEASUREMENT,
+    SensorDeviceClass,
+    SensorStateClass,
 )
 from homeassistant.const import (
     ATTR_DEVICE_CLASS,
@@ -18,7 +19,6 @@ from homeassistant.const import (
     CONCENTRATION_PARTS_PER_BILLION,
     CONCENTRATION_PARTS_PER_CUBIC_METER,
     CONCENTRATION_PARTS_PER_MILLION,
-    DEVICE_CLASS_CO,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
@@ -42,7 +42,7 @@ async def test_air_quality(
     assert entry.unique_id == f"{entry_id}_air_quality_particulate_matter_2_5"
     assert state.state == "3.14"
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Particulate Matter < 2.5 μm"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_MICROGRAMS_PER_CUBIC_METER
@@ -57,7 +57,7 @@ async def test_air_quality(
     assert entry.unique_id == f"{entry_id}_air_quality_particulate_matter_10"
     assert state.state == "5.24"
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Particulate Matter < 10 μm"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_MICROGRAMS_PER_CUBIC_METER
@@ -72,7 +72,7 @@ async def test_air_quality(
     assert entry.unique_id == f"{entry_id}_air_quality_sulphur_dioxide"
     assert state.state == "0.031"
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Sulphur Dioxide (SO2)"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_PARTS_PER_BILLION
@@ -87,7 +87,7 @@ async def test_air_quality(
     assert entry.unique_id == f"{entry_id}_air_quality_nitrogen_dioxide"
     assert state.state == "0.66"
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Nitrogen Dioxide (NO2)"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_PARTS_PER_BILLION
@@ -102,7 +102,7 @@ async def test_air_quality(
     assert entry.unique_id == f"{entry_id}_air_quality_ozone"
     assert state.state == "17.067"
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Ozone"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_PARTS_PER_BILLION
@@ -116,9 +116,9 @@ async def test_air_quality(
     assert state
     assert entry.unique_id == f"{entry_id}_air_quality_carbon_monoxide"
     assert state.state == "0.105"
-    assert state.attributes.get(ATTR_DEVICE_CLASS) == DEVICE_CLASS_CO
+    assert state.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.CO
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Carbon Monoxide (CO)"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_PARTS_PER_MILLION
@@ -132,7 +132,7 @@ async def test_air_quality(
     assert entry.unique_id == f"{entry_id}_air_quality_air_quality_index"
     assert state.state == "13"
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Air Quality Index (AQI)"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert ATTR_DEVICE_CLASS not in state.attributes
     assert ATTR_UNIT_OF_MEASUREMENT not in state.attributes
     assert ATTR_ICON not in state.attributes
@@ -165,7 +165,7 @@ async def test_pollen(
     assert state.state == "190"
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Grass Pollen"
     assert state.attributes.get(ATTR_ICON) == "mdi:grass"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_PARTS_PER_CUBIC_METER
@@ -180,7 +180,7 @@ async def test_pollen(
     assert state.state == "127"
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Tree Pollen"
     assert state.attributes.get(ATTR_ICON) == "mdi:tree"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_PARTS_PER_CUBIC_METER
@@ -195,7 +195,7 @@ async def test_pollen(
     assert state.state == "95"
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == "Weed Pollen"
     assert state.attributes.get(ATTR_ICON) == "mdi:sprout"
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_PARTS_PER_CUBIC_METER
@@ -337,7 +337,7 @@ async def test_pollen_enable_disable_by_defaults(
     assert state.state == value
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == name
     assert state.attributes.get(ATTR_ICON) == icon
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
         == CONCENTRATION_PARTS_PER_CUBIC_METER