From 997a847b5cb4882dfca9ca250e9fb1f7bdaafdd7 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Wed, 26 May 2021 00:21:18 +0200
Subject: [PATCH] Add support for Sensor state class to ESPHome (#51090)

* Add support for Sensor state class to ESPHome

* Bump aioesphome to 2.8.0
---
 .../components/esphome/manifest.json          |  2 +-
 homeassistant/components/esphome/sensor.py    | 31 +++++++++++++++++--
 requirements_all.txt                          |  2 +-
 requirements_test_all.txt                     |  2 +-
 4 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/esphome/manifest.json b/homeassistant/components/esphome/manifest.json
index e103fa65992..592ca616d04 100644
--- a/homeassistant/components/esphome/manifest.json
+++ b/homeassistant/components/esphome/manifest.json
@@ -3,7 +3,7 @@
   "name": "ESPHome",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/esphome",
-  "requirements": ["aioesphomeapi==2.7.0"],
+  "requirements": ["aioesphomeapi==2.8.0"],
   "zeroconf": ["_esphomelib._tcp.local."],
   "codeowners": ["@OttoWinter"],
   "after_dependencies": ["zeroconf", "tag"],
diff --git a/homeassistant/components/esphome/sensor.py b/homeassistant/components/esphome/sensor.py
index 12319be8c40..7b905aad148 100644
--- a/homeassistant/components/esphome/sensor.py
+++ b/homeassistant/components/esphome/sensor.py
@@ -3,12 +3,19 @@ from __future__ import annotations
 
 import math
 
-from aioesphomeapi import SensorInfo, SensorState, TextSensorInfo, TextSensorState
+from aioesphomeapi import (
+    SensorInfo,
+    SensorState,
+    SensorStateClass,
+    TextSensorInfo,
+    TextSensorState,
+)
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
     DEVICE_CLASS_TIMESTAMP,
     DEVICE_CLASSES,
+    STATE_CLASS_MEASUREMENT,
     SensorEntity,
 )
 from homeassistant.config_entries import ConfigEntry
@@ -16,7 +23,12 @@ from homeassistant.core import HomeAssistant
 import homeassistant.helpers.config_validation as cv
 from homeassistant.util import dt
 
-from . import EsphomeEntity, esphome_state_property, platform_async_setup_entry
+from . import (
+    EsphomeEntity,
+    esphome_map_enum,
+    esphome_state_property,
+    platform_async_setup_entry,
+)
 
 ICON_SCHEMA = vol.Schema(cv.icon)
 
@@ -49,6 +61,14 @@ async def async_setup_entry(
 # pylint: disable=invalid-overridden-method
 
 
+@esphome_map_enum
+def _state_classes():
+    return {
+        SensorStateClass.NONE: None,
+        SensorStateClass.MEASUREMENT: STATE_CLASS_MEASUREMENT,
+    }
+
+
 class EsphomeSensor(EsphomeEntity, SensorEntity):
     """A sensor implementation for esphome."""
 
@@ -97,6 +117,13 @@ class EsphomeSensor(EsphomeEntity, SensorEntity):
             return None
         return self._static_info.device_class
 
+    @property
+    def state_class(self) -> str | None:
+        """Return the state class of this entity."""
+        if not self._static_info.state_class:
+            return None
+        return _state_classes.from_esphome(self._static_info.state_class)
+
 
 class EsphomeTextSensor(EsphomeEntity, SensorEntity):
     """A text sensor implementation for ESPHome."""
diff --git a/requirements_all.txt b/requirements_all.txt
index 0076da48f34..a3de0bb383e 100644
--- a/requirements_all.txt
+++ b/requirements_all.txt
@@ -160,7 +160,7 @@ aioeafm==0.1.2
 aioemonitor==1.0.5
 
 # homeassistant.components.esphome
-aioesphomeapi==2.7.0
+aioesphomeapi==2.8.0
 
 # homeassistant.components.flo
 aioflo==0.4.1
diff --git a/requirements_test_all.txt b/requirements_test_all.txt
index 73831e1f90b..015e565af0a 100644
--- a/requirements_test_all.txt
+++ b/requirements_test_all.txt
@@ -100,7 +100,7 @@ aioeafm==0.1.2
 aioemonitor==1.0.5
 
 # homeassistant.components.esphome
-aioesphomeapi==2.7.0
+aioesphomeapi==2.8.0
 
 # homeassistant.components.flo
 aioflo==0.4.1