From ba7ad8a58fd71e648917084a769543f5b5919b42 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Hjelseth=20H=C3=B8yer?= <mail@dahoiv.net>
Date: Tue, 29 Jun 2021 19:16:43 +0200
Subject: [PATCH] Add Melcloud device class and state class (#52276)

---
 homeassistant/components/melcloud/sensor.py   | 46 +++++++------------
 .../melcloud/test_atw_zone_sensor.py          | 16 +++++--
 2 files changed, 30 insertions(+), 32 deletions(-)

diff --git a/homeassistant/components/melcloud/sensor.py b/homeassistant/components/melcloud/sensor.py
index 2c59763dc72..c1b7e5e8cbd 100644
--- a/homeassistant/components/melcloud/sensor.py
+++ b/homeassistant/components/melcloud/sensor.py
@@ -2,14 +2,19 @@
 from pymelcloud import DEVICE_TYPE_ATA, DEVICE_TYPE_ATW
 from pymelcloud.atw_device import Zone
 
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import (
+    DEVICE_CLASS_ENERGY,
+    DEVICE_CLASS_TEMPERATURE,
+    STATE_CLASS_MEASUREMENT,
+    SensorEntity,
+)
 from homeassistant.const import (
     ATTR_DEVICE_CLASS,
     ATTR_ICON,
-    DEVICE_CLASS_TEMPERATURE,
     ENERGY_KILO_WATT_HOUR,
     TEMP_CELSIUS,
 )
+from homeassistant.util import dt as dt_util
 
 from . import MelCloudDevice
 from .const import DOMAIN
@@ -32,7 +37,7 @@ ATA_SENSORS = {
         ATTR_MEASUREMENT_NAME: "Energy",
         ATTR_ICON: "mdi:factory",
         ATTR_UNIT: ENERGY_KILO_WATT_HOUR,
-        ATTR_DEVICE_CLASS: None,
+        ATTR_DEVICE_CLASS: DEVICE_CLASS_ENERGY,
         ATTR_VALUE_FN: lambda x: x.device.total_energy_consumed,
         ATTR_ENABLED_FN: lambda x: x.device.has_energy_consumed_meter,
     },
@@ -116,40 +121,23 @@ class MelDeviceSensor(SensorEntity):
     def __init__(self, api: MelCloudDevice, measurement, definition):
         """Initialize the sensor."""
         self._api = api
-        self._name_slug = api.name
-        self._measurement = measurement
         self._def = definition
 
-    @property
-    def unique_id(self):
-        """Return a unique ID."""
-        return f"{self._api.device.serial}-{self._api.device.mac}-{self._measurement}"
+        self._attr_device_class = definition[ATTR_DEVICE_CLASS]
+        self._attr_icon = definition[ATTR_ICON]
+        self._attr_name = f"{api.name} {definition[ATTR_MEASUREMENT_NAME]}"
+        self._attr_unique_id = f"{api.device.serial}-{api.device.mac}-{measurement}"
+        self._attr_unit_of_measurement = definition[ATTR_UNIT]
+        self._attr_state_class = STATE_CLASS_MEASUREMENT
 
-    @property
-    def icon(self):
-        """Return the icon to use in the frontend, if any."""
-        return self._def[ATTR_ICON]
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return f"{self._name_slug} {self._def[ATTR_MEASUREMENT_NAME]}"
+        if self.device_class == DEVICE_CLASS_ENERGY:
+            self._attr_last_reset = dt_util.utc_from_timestamp(0)
 
     @property
     def state(self):
         """Return the state of the sensor."""
         return self._def[ATTR_VALUE_FN](self._api)
 
-    @property
-    def unit_of_measurement(self):
-        """Return the unit of measurement."""
-        return self._def[ATTR_UNIT]
-
-    @property
-    def device_class(self):
-        """Return device class."""
-        return self._def[ATTR_DEVICE_CLASS]
-
     async def async_update(self):
         """Retrieve latest state."""
         await self._api.async_update()
@@ -171,7 +159,7 @@ class AtwZoneSensor(MelDeviceSensor):
             full_measurement = f"{measurement}-zone-{zone.zone_index}"
         super().__init__(api, full_measurement, definition)
         self._zone = zone
-        self._name_slug = f"{api.name} {zone.name}"
+        self._attr_name = f"{api.name} {zone.name} {definition[ATTR_MEASUREMENT_NAME]}"
 
     @property
     def state(self):
diff --git a/tests/components/melcloud/test_atw_zone_sensor.py b/tests/components/melcloud/test_atw_zone_sensor.py
index ac34a5ccc49..6e6487a3774 100644
--- a/tests/components/melcloud/test_atw_zone_sensor.py
+++ b/tests/components/melcloud/test_atw_zone_sensor.py
@@ -3,7 +3,7 @@ from unittest.mock import patch
 
 import pytest
 
-from homeassistant.components.melcloud.sensor import AtwZoneSensor
+from homeassistant.components.melcloud.sensor import ATW_ZONE_SENSORS, AtwZoneSensor
 
 
 @pytest.fixture
@@ -34,8 +34,18 @@ def mock_zone_2():
 
 def test_zone_unique_ids(mock_device, mock_zone_1, mock_zone_2):
     """Test unique id generation correctness."""
-    sensor_1 = AtwZoneSensor(mock_device, mock_zone_1, "room_temperature", {})
+    sensor_1 = AtwZoneSensor(
+        mock_device,
+        mock_zone_1,
+        "room_temperature",
+        ATW_ZONE_SENSORS["room_temperature"],
+    )
     assert sensor_1.unique_id == "1234-11:11:11:11:11:11-room_temperature"
 
-    sensor_2 = AtwZoneSensor(mock_device, mock_zone_2, "room_temperature", {})
+    sensor_2 = AtwZoneSensor(
+        mock_device,
+        mock_zone_2,
+        "room_temperature",
+        ATW_ZONE_SENSORS["flow_temperature"],
+    )
     assert sensor_2.unique_id == "1234-11:11:11:11:11:11-room_temperature-zone-2"