From 2b0b8736f2d181159ae7a32f2681bdd3c2daa4d6 Mon Sep 17 00:00:00 2001
From: Michael <35783820+mib1185@users.noreply.github.com>
Date: Thu, 29 Jul 2021 20:54:51 +0200
Subject: [PATCH] Set state class measurement also for Total Energy for AVM
 Fritz!Smarthome devices (#53707)

---
 homeassistant/components/fritzbox/sensor.py | 2 +-
 tests/components/fritzbox/test_switch.py    | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/fritzbox/sensor.py b/homeassistant/components/fritzbox/sensor.py
index 24b3d9cc5ff..d325e592faf 100644
--- a/homeassistant/components/fritzbox/sensor.py
+++ b/homeassistant/components/fritzbox/sensor.py
@@ -96,7 +96,7 @@ async def async_setup_entry(
                         ATTR_ENTITY_ID: f"{device.ain}_total_energy",
                         ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR,
                         ATTR_DEVICE_CLASS: DEVICE_CLASS_ENERGY,
-                        ATTR_STATE_CLASS: None,
+                        ATTR_STATE_CLASS: STATE_CLASS_MEASUREMENT,
                     },
                     coordinator,
                     ain,
diff --git a/tests/components/fritzbox/test_switch.py b/tests/components/fritzbox/test_switch.py
index cb6ae85f889..951528f1e7d 100644
--- a/tests/components/fritzbox/test_switch.py
+++ b/tests/components/fritzbox/test_switch.py
@@ -76,7 +76,7 @@ async def test_setup(hass: HomeAssistant, fritz: Mock):
     assert state.attributes[ATTR_LAST_RESET] == "1970-01-01T00:00:00+00:00"
     assert state.attributes[ATTR_FRIENDLY_NAME] == f"{CONF_FAKE_NAME} Total Energy"
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == ENERGY_KILO_WATT_HOUR
-    assert ATTR_STATE_CLASS not in state.attributes
+    assert state.attributes[ATTR_STATE_CLASS] == STATE_CLASS_MEASUREMENT
 
 
 async def test_turn_on(hass: HomeAssistant, fritz: Mock):