From 3bfcca2bb06c2d8b51e53cf240fb8eb3633d375d Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Wed, 23 Jun 2021 19:23:56 +0200
Subject: [PATCH] Add state class to Atome Linky, use class attributes (#52107)

---
 homeassistant/components/atome/sensor.py | 51 +++++++++---------------
 1 file changed, 18 insertions(+), 33 deletions(-)

diff --git a/homeassistant/components/atome/sensor.py b/homeassistant/components/atome/sensor.py
index f47dd4033b9..bcb7b4f1ece 100644
--- a/homeassistant/components/atome/sensor.py
+++ b/homeassistant/components/atome/sensor.py
@@ -5,7 +5,11 @@ import logging
 from pyatome.client import AtomeClient, PyAtomeError
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA,
+    STATE_CLASS_MEASUREMENT,
+    SensorEntity,
+)
 from homeassistant.const import (
     CONF_NAME,
     CONF_PASSWORD,
@@ -219,37 +223,16 @@ class AtomeSensor(SensorEntity):
 
     def __init__(self, data, name, sensor_type):
         """Initialize the sensor."""
-        self._name = name
+        self._attr_name = name
         self._data = data
-        self._state = None
-        self._attributes = {}
 
         self._sensor_type = sensor_type
 
         if sensor_type == LIVE_TYPE:
-            self._unit_of_measurement = POWER_WATT
+            self._attr_unit_of_measurement = POWER_WATT
+            self._attr_state_class = STATE_CLASS_MEASUREMENT
         else:
-            self._unit_of_measurement = ENERGY_KILO_WATT_HOUR
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return self._name
-
-    @property
-    def state(self):
-        """Return the state of the sensor."""
-        return self._state
-
-    @property
-    def extra_state_attributes(self):
-        """Return the state attributes."""
-        return self._attributes
-
-    @property
-    def unit_of_measurement(self):
-        """Return the unit of measurement."""
-        return self._unit_of_measurement
+            self._attr_unit_of_measurement = ENERGY_KILO_WATT_HOUR
 
     def update(self):
         """Update device state."""
@@ -257,11 +240,13 @@ class AtomeSensor(SensorEntity):
         update_function()
 
         if self._sensor_type == LIVE_TYPE:
-            self._state = self._data.live_power
-            self._attributes["subscribed_power"] = self._data.subscribed_power
-            self._attributes["is_connected"] = self._data.is_connected
+            self._attr_state = self._data.live_power
+            self._attr_extra_state_attributes = {
+                "subscribed_power": self._data.subscribed_power,
+                "is_connected": self._data.is_connected,
+            }
         else:
-            self._state = getattr(self._data, f"{self._sensor_type}_usage")
-            self._attributes["price"] = getattr(
-                self._data, f"{self._sensor_type}_price"
-            )
+            self._attr_state = getattr(self._data, f"{self._sensor_type}_usage")
+            self._attr_extra_state_attributes = {
+                "price": getattr(self._data, f"{self._sensor_type}_price")
+            }