Support device and state classes for WAQI sensor (#57762)

homeassistant/components/waqi/sensor.py

@@ -7,12 +7,13 @@ import aiohttp
 import voluptuous as vol
 from waqiasync import WaqiClient
 
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import STATE_CLASS_MEASUREMENT, SensorEntity
 from homeassistant.const import (
     ATTR_ATTRIBUTION,
     ATTR_TEMPERATURE,
     ATTR_TIME,
     CONF_TOKEN,
+    DEVICE_CLASS_AQI,
 )
 from homeassistant.exceptions import PlatformNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -43,6 +44,9 @@ KEY_TO_ATTR = {
 
 ATTRIBUTION = "Data provided by the World Air Quality Index project"
 
+ATTR_ICON = "mdi:cloud"
+ATTR_UNIT = "AQI"
+
 CONF_LOCATIONS = "locations"
 CONF_STATIONS = "stations"
 
@@ -96,6 +100,11 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
 class WaqiSensor(SensorEntity):
     """Implementation of a WAQI sensor."""
 
+    _attr_icon = ATTR_ICON
+    _attr_native_unit_of_measurement = ATTR_UNIT
+    _attr_device_class = DEVICE_CLASS_AQI
+    _attr_state_class = STATE_CLASS_MEASUREMENT
+
     def __init__(self, client, station):
         """Initialize the sensor."""
         self._client = client
@@ -123,11 +132,6 @@ class WaqiSensor(SensorEntity):
             return f"WAQI {self.station_name}"
         return f"WAQI {self.url if self.url else self.uid}"
 
-    @property
-    def icon(self):
-        """Icon to use in the frontend, if any."""
-        return "mdi:cloud"
-
     @property
     def native_value(self):
         """Return the state of the device."""
@@ -145,11 +149,6 @@ class WaqiSensor(SensorEntity):
         """Return unique ID."""
         return self.uid
 
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit of measurement of this entity, if any."""
-        return "AQI"
-
     @property
     def extra_state_attributes(self):
         """Return the state attributes of the last update."""