From 9d84d41f815c58bf15bdea76fc78e7ccc8249f7b Mon Sep 17 00:00:00 2001
From: Jean-Yves Avenard <jyavenard@gmail.com>
Date: Wed, 6 Oct 2021 22:24:17 +1100
Subject: [PATCH] Change energy state class to STATE_CLASS_TOTAL (#56974)

---
 homeassistant/components/iotawatt/sensor.py |  4 ++--
 tests/components/iotawatt/test_sensor.py    | 10 +++++-----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/iotawatt/sensor.py b/homeassistant/components/iotawatt/sensor.py
index ba0ec30caa0..ec2918b0ce6 100644
--- a/homeassistant/components/iotawatt/sensor.py
+++ b/homeassistant/components/iotawatt/sensor.py
@@ -9,7 +9,7 @@ from iotawattpy.sensor import Sensor
 
 from homeassistant.components.sensor import (
     STATE_CLASS_MEASUREMENT,
-    STATE_CLASS_TOTAL_INCREASING,
+    STATE_CLASS_TOTAL,
     SensorEntity,
     SensorEntityDescription,
 )
@@ -83,6 +83,7 @@ ENTITY_DESCRIPTION_KEY_MAP: dict[str, IotaWattSensorEntityDescription] = {
     "WattHours": IotaWattSensorEntityDescription(
         "WattHours",
         native_unit_of_measurement=ENERGY_WATT_HOUR,
+        state_class=STATE_CLASS_TOTAL,
         device_class=DEVICE_CLASS_ENERGY,
     ),
     "VA": IotaWattSensorEntityDescription(
@@ -242,7 +243,6 @@ class IotaWattAccumulatingSensor(IotaWattSensor, RestoreEntity):
 
         super().__init__(coordinator, key, entity_description)
 
-        self._attr_state_class = STATE_CLASS_TOTAL_INCREASING
         if self._attr_unique_id is not None:
             self._attr_unique_id += ".accumulated"
 
diff --git a/tests/components/iotawatt/test_sensor.py b/tests/components/iotawatt/test_sensor.py
index 8928c012d48..2397338c22c 100644
--- a/tests/components/iotawatt/test_sensor.py
+++ b/tests/components/iotawatt/test_sensor.py
@@ -6,7 +6,7 @@ from homeassistant.components.sensor import (
     ATTR_STATE_CLASS,
     DEVICE_CLASS_ENERGY,
     STATE_CLASS_MEASUREMENT,
-    STATE_CLASS_TOTAL_INCREASING,
+    STATE_CLASS_TOTAL,
 )
 from homeassistant.const import (
     ATTR_DEVICE_CLASS,
@@ -74,7 +74,7 @@ async def test_sensor_type_output(hass, mock_iotawatt):
     state = hass.states.get("sensor.my_watthour_sensor")
     assert state is not None
     assert state.state == "243"
-    assert ATTR_STATE_CLASS not in state.attributes
+    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_TOTAL
     assert state.attributes[ATTR_FRIENDLY_NAME] == "My WattHour Sensor"
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == ENERGY_WATT_HOUR
     assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_ENERGY
@@ -125,7 +125,7 @@ async def test_sensor_type_accumulated_output(hass, mock_iotawatt):
         state.attributes[ATTR_FRIENDLY_NAME]
         == "My WattHour Accumulated Output Sensor.wh Accumulated"
     )
-    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_TOTAL_INCREASING
+    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_TOTAL
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == ENERGY_WATT_HOUR
     assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_ENERGY
     assert state.attributes["type"] == "Output"
@@ -166,7 +166,7 @@ async def test_sensor_type_accumulated_output_error_restore(hass, mock_iotawatt)
         state.attributes[ATTR_FRIENDLY_NAME]
         == "My WattHour Accumulated Output Sensor.wh Accumulated"
     )
-    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_TOTAL_INCREASING
+    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_TOTAL
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == ENERGY_WATT_HOUR
     assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_ENERGY
     assert state.attributes["type"] == "Output"
@@ -224,7 +224,7 @@ async def test_sensor_type_multiple_accumulated_output(hass, mock_iotawatt):
         state.attributes[ATTR_FRIENDLY_NAME]
         == "My WattHour Accumulated Output Sensor.wh Accumulated"
     )
-    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_TOTAL_INCREASING
+    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_TOTAL
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == ENERGY_WATT_HOUR
     assert state.attributes[ATTR_DEVICE_CLASS] == DEVICE_CLASS_ENERGY
     assert state.attributes["type"] == "Output"