From 33bddf275c9d2f33bcff82fa93292b132e85c988 Mon Sep 17 00:00:00 2001
From: Dave T <17680170+davet2001@users.noreply.github.com>
Date: Tue, 21 Dec 2021 18:17:23 +0000
Subject: [PATCH] Use SensorStateClass Enums in fritzbox tests (#62134)

---
 tests/components/fritzbox/test_binary_sensor.py |  4 ++--
 tests/components/fritzbox/test_sensor.py        | 10 +++-------
 tests/components/fritzbox/test_switch.py        |  9 ++++-----
 3 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/tests/components/fritzbox/test_binary_sensor.py b/tests/components/fritzbox/test_binary_sensor.py
index f010c8499ea..a1ccc2bd0f2 100644
--- a/tests/components/fritzbox/test_binary_sensor.py
+++ b/tests/components/fritzbox/test_binary_sensor.py
@@ -5,7 +5,7 @@ from unittest.mock import Mock
 
 from requests.exceptions import HTTPError
 
-from homeassistant.components.binary_sensor import DOMAIN
+from homeassistant.components.binary_sensor import DOMAIN, BinarySensorDeviceClass
 from homeassistant.components.fritzbox.const import DOMAIN as FB_DOMAIN
 from homeassistant.components.sensor import ATTR_STATE_CLASS, DOMAIN as SENSOR_DOMAIN
 from homeassistant.const import (
@@ -39,7 +39,7 @@ async def test_setup(hass: HomeAssistant, fritz: Mock):
     assert state
     assert state.state == STATE_ON
     assert state.attributes[ATTR_FRIENDLY_NAME] == CONF_FAKE_NAME
-    assert state.attributes[ATTR_DEVICE_CLASS] == "window"
+    assert state.attributes[ATTR_DEVICE_CLASS] == BinarySensorDeviceClass.WINDOW
     assert ATTR_STATE_CLASS not in state.attributes
 
     state = hass.states.get(f"{SENSOR_DOMAIN}.{CONF_FAKE_NAME}_battery")
diff --git a/tests/components/fritzbox/test_sensor.py b/tests/components/fritzbox/test_sensor.py
index 42cf90bba58..07b54765c9a 100644
--- a/tests/components/fritzbox/test_sensor.py
+++ b/tests/components/fritzbox/test_sensor.py
@@ -9,11 +9,7 @@ from homeassistant.components.fritzbox.const import (
     ATTR_STATE_LOCKED,
     DOMAIN as FB_DOMAIN,
 )
-from homeassistant.components.sensor import (
-    ATTR_STATE_CLASS,
-    DOMAIN,
-    STATE_CLASS_MEASUREMENT,
-)
+from homeassistant.components.sensor import ATTR_STATE_CLASS, DOMAIN, SensorStateClass
 from homeassistant.const import (
     ATTR_FRIENDLY_NAME,
     ATTR_UNIT_OF_MEASUREMENT,
@@ -47,14 +43,14 @@ async def test_setup(hass: HomeAssistant, fritz: Mock):
     assert state.attributes[ATTR_STATE_DEVICE_LOCKED] == "fake_locked_device"
     assert state.attributes[ATTR_STATE_LOCKED] == "fake_locked"
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == TEMP_CELSIUS
-    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_MEASUREMENT
+    assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
 
     state = hass.states.get(f"{ENTITY_ID}_humidity")
     assert state
     assert state.state == "42"
     assert state.attributes[ATTR_FRIENDLY_NAME] == f"{CONF_FAKE_NAME} Humidity"
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == PERCENTAGE
-    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_MEASUREMENT
+    assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
 
     state = hass.states.get(f"{ENTITY_ID}_battery")
     assert state
diff --git a/tests/components/fritzbox/test_switch.py b/tests/components/fritzbox/test_switch.py
index 73bb7a1110b..f6f81d5f24c 100644
--- a/tests/components/fritzbox/test_switch.py
+++ b/tests/components/fritzbox/test_switch.py
@@ -12,8 +12,7 @@ from homeassistant.components.fritzbox.const import (
 from homeassistant.components.sensor import (
     ATTR_STATE_CLASS,
     DOMAIN as SENSOR_DOMAIN,
-    STATE_CLASS_MEASUREMENT,
-    STATE_CLASS_TOTAL_INCREASING,
+    SensorStateClass,
 )
 from homeassistant.components.switch import DOMAIN
 from homeassistant.const import (
@@ -62,7 +61,7 @@ async def test_setup(hass: HomeAssistant, fritz: Mock):
     assert state.attributes[ATTR_STATE_DEVICE_LOCKED] == "fake_locked_device"
     assert state.attributes[ATTR_STATE_LOCKED] == "fake_locked"
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == TEMP_CELSIUS
-    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_MEASUREMENT
+    assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
 
     state = hass.states.get(f"{ENTITY_ID}_humidity")
     assert state is None
@@ -72,14 +71,14 @@ async def test_setup(hass: HomeAssistant, fritz: Mock):
     assert state.state == "5.678"
     assert state.attributes[ATTR_FRIENDLY_NAME] == f"{CONF_FAKE_NAME} Power Consumption"
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == POWER_WATT
-    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_MEASUREMENT
+    assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.MEASUREMENT
 
     state = hass.states.get(f"{SENSOR_DOMAIN}.{CONF_FAKE_NAME}_total_energy")
     assert state
     assert state.state == "1.234"
     assert state.attributes[ATTR_FRIENDLY_NAME] == f"{CONF_FAKE_NAME} Total Energy"
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == ENERGY_KILO_WATT_HOUR
-    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_TOTAL_INCREASING
+    assert state.attributes[ATTR_STATE_CLASS] == SensorStateClass.TOTAL_INCREASING
 
 
 async def test_turn_on(hass: HomeAssistant, fritz: Mock):