From af631b90e5f47c5b6d10d156f2cd46ceee612def Mon Sep 17 00:00:00 2001
From: Robert Hillis <tkdrob4390@yahoo.com>
Date: Sat, 18 Dec 2021 08:54:26 -0500
Subject: [PATCH] Use enums in statistics tests (#62191)

---
 homeassistant/components/statistics/sensor.py | 4 ++--
 tests/components/statistics/test_sensor.py    | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/statistics/sensor.py b/homeassistant/components/statistics/sensor.py
index b06589f52c0..e2608c888aa 100644
--- a/homeassistant/components/statistics/sensor.py
+++ b/homeassistant/components/statistics/sensor.py
@@ -10,8 +10,8 @@ from homeassistant.components.recorder.models import States
 from homeassistant.components.recorder.util import execute, session_scope
 from homeassistant.components.sensor import (
     PLATFORM_SCHEMA,
-    STATE_CLASS_MEASUREMENT,
     SensorEntity,
+    SensorStateClass,
 )
 from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
@@ -350,7 +350,7 @@ class StatisticsSensor(SensorEntity):
         """Return the state class of this entity."""
         if self._state_characteristic in STATS_NOT_A_NUMBER:
             return None
-        return STATE_CLASS_MEASUREMENT
+        return SensorStateClass.MEASUREMENT
 
     @property
     def native_value(self):
diff --git a/tests/components/statistics/test_sensor.py b/tests/components/statistics/test_sensor.py
index e3ace35f0ae..51ab31e6ed5 100644
--- a/tests/components/statistics/test_sensor.py
+++ b/tests/components/statistics/test_sensor.py
@@ -4,7 +4,7 @@ import statistics
 from unittest.mock import patch
 
 from homeassistant import config as hass_config
-from homeassistant.components.sensor import ATTR_STATE_CLASS, STATE_CLASS_MEASUREMENT
+from homeassistant.components.sensor import ATTR_STATE_CLASS, SensorStateClass
 from homeassistant.components.statistics import DOMAIN as STATISTICS_DOMAIN
 from homeassistant.components.statistics.sensor import StatisticsSensor
 from homeassistant.const import (
@@ -82,7 +82,7 @@ async def test_sensor_defaults_numeric(hass):
     state = hass.states.get("sensor.test")
     assert state.state == str(round(sum(VALUES_NUMERIC) / len(VALUES_NUMERIC), 2))
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == TEMP_CELSIUS
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) is SensorStateClass.MEASUREMENT
     assert state.attributes.get("buffer_usage_ratio") == round(9 / 20, 2)
     assert state.attributes.get("source_value_valid") is True
     assert "age_coverage_ratio" not in state.attributes
@@ -167,7 +167,7 @@ async def test_sensor_defaults_binary(hass):
     state = hass.states.get("sensor.test")
     assert state.state == str(len(VALUES_BINARY))
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) is None
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) is SensorStateClass.MEASUREMENT
     assert state.attributes.get("buffer_usage_ratio") == round(9 / 20, 2)
     assert state.attributes.get("source_value_valid") is True
     assert "age_coverage_ratio" not in state.attributes
@@ -444,7 +444,7 @@ async def test_state_class(hass):
     await hass.async_block_till_done()
 
     state = hass.states.get("sensor.test_normal")
-    assert state.attributes.get(ATTR_STATE_CLASS) == STATE_CLASS_MEASUREMENT
+    assert state.attributes.get(ATTR_STATE_CLASS) is SensorStateClass.MEASUREMENT
     state = hass.states.get("sensor.test_nan")
     assert state.attributes.get(ATTR_STATE_CLASS) is None