From 056e71266763446b7043d1bb9fdd8821d10de8c5 Mon Sep 17 00:00:00 2001
From: Chris Talkington <chris@talkingtontech.com>
Date: Sun, 13 Sep 2020 21:29:59 -0500
Subject: [PATCH] Add device class to canary sensors (#40050)

* add device class to canary sensors

* Update test_sensor.py

* Update sensor.py

* Update sensor.py
---
 homeassistant/components/canary/sensor.py | 32 ++++++++++++++---------
 tests/components/canary/test_sensor.py    | 24 +++++++++++++----
 2 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/homeassistant/components/canary/sensor.py b/homeassistant/components/canary/sensor.py
index 5f1b1fe906b..1af2b5ad135 100644
--- a/homeassistant/components/canary/sensor.py
+++ b/homeassistant/components/canary/sensor.py
@@ -1,9 +1,15 @@
 """Support for Canary sensors."""
 from canary.api import SensorType
 
-from homeassistant.const import PERCENTAGE, TEMP_CELSIUS
+from homeassistant.const import (
+    DEVICE_CLASS_BATTERY,
+    DEVICE_CLASS_HUMIDITY,
+    DEVICE_CLASS_SIGNAL_STRENGTH,
+    DEVICE_CLASS_TEMPERATURE,
+    PERCENTAGE,
+    TEMP_CELSIUS,
+)
 from homeassistant.helpers.entity import Entity
-from homeassistant.helpers.icon import icon_for_battery_level
 
 from . import DATA_CANARY
 
@@ -18,13 +24,13 @@ CANARY_PRO = "Canary Pro"
 CANARY_FLEX = "Canary Flex"
 
 # Sensor types are defined like so:
-# sensor type name, unit_of_measurement, icon
+# sensor type name, unit_of_measurement, icon, device class, products supported
 SENSOR_TYPES = [
-    ["temperature", TEMP_CELSIUS, "mdi:thermometer", [CANARY_PRO]],
-    ["humidity", PERCENTAGE, "mdi:water-percent", [CANARY_PRO]],
-    ["air_quality", None, "mdi:weather-windy", [CANARY_PRO]],
-    ["wifi", "dBm", "mdi:wifi", [CANARY_FLEX]],
-    ["battery", PERCENTAGE, "mdi:battery-50", [CANARY_FLEX]],
+    ["temperature", TEMP_CELSIUS, None, DEVICE_CLASS_TEMPERATURE, [CANARY_PRO]],
+    ["humidity", PERCENTAGE, None, DEVICE_CLASS_HUMIDITY, [CANARY_PRO]],
+    ["air_quality", None, "mdi:weather-windy", None, [CANARY_PRO]],
+    ["wifi", "dBm", None, DEVICE_CLASS_SIGNAL_STRENGTH, [CANARY_FLEX]],
+    ["battery", PERCENTAGE, None, DEVICE_CLASS_BATTERY, [CANARY_FLEX]],
 ]
 
 STATE_AIR_QUALITY_NORMAL = "normal"
@@ -42,7 +48,7 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
             if device.is_online:
                 device_type = device.device_type
                 for sensor_type in SENSOR_TYPES:
-                    if device_type.get("name") in sensor_type[3]:
+                    if device_type.get("name") in sensor_type[4]:
                         devices.append(
                             CanarySensor(data, sensor_type, location, device)
                         )
@@ -83,12 +89,14 @@ class CanarySensor(Entity):
         """Return the unit of measurement."""
         return self._sensor_type[1]
 
+    @property
+    def device_class(self):
+        """Device class for the sensor."""
+        return self._sensor_type[3]
+
     @property
     def icon(self):
         """Icon for the sensor."""
-        if self.state is not None and self._sensor_type[0] == "battery":
-            return icon_for_battery_level(battery_level=self.state)
-
         return self._sensor_type[2]
 
     @property
diff --git a/tests/components/canary/test_sensor.py b/tests/components/canary/test_sensor.py
index 02d2dc5cc24..8d785a6ced5 100644
--- a/tests/components/canary/test_sensor.py
+++ b/tests/components/canary/test_sensor.py
@@ -6,7 +6,15 @@ from homeassistant.components.canary.sensor import (
     STATE_AIR_QUALITY_NORMAL,
     STATE_AIR_QUALITY_VERY_ABNORMAL,
 )
-from homeassistant.const import ATTR_UNIT_OF_MEASUREMENT, PERCENTAGE, TEMP_CELSIUS
+from homeassistant.const import (
+    ATTR_UNIT_OF_MEASUREMENT,
+    DEVICE_CLASS_BATTERY,
+    DEVICE_CLASS_HUMIDITY,
+    DEVICE_CLASS_SIGNAL_STRENGTH,
+    DEVICE_CLASS_TEMPERATURE,
+    PERCENTAGE,
+    TEMP_CELSIUS,
+)
 from homeassistant.setup import async_setup_component
 
 from . import mock_device, mock_location, mock_reading
@@ -43,15 +51,15 @@ async def test_sensors_pro(hass, canary) -> None:
             "20_temperature",
             "21.12",
             TEMP_CELSIUS,
+            DEVICE_CLASS_TEMPERATURE,
             None,
-            "mdi:thermometer",
         ),
         "home_dining_room_humidity": (
             "20_humidity",
             "50.46",
             PERCENTAGE,
+            DEVICE_CLASS_HUMIDITY,
             None,
-            "mdi:water-percent",
         ),
         "home_dining_room_air_quality": (
             "20_air_quality",
@@ -156,10 +164,16 @@ async def test_sensors_flex(hass, canary) -> None:
             "20_battery",
             "70.46",
             PERCENTAGE,
+            DEVICE_CLASS_BATTERY,
+            None,
+        ),
+        "home_dining_room_wifi": (
+            "20_wifi",
+            "-57.0",
+            "dBm",
+            DEVICE_CLASS_SIGNAL_STRENGTH,
             None,
-            "mdi:battery-70",
         ),
-        "home_dining_room_wifi": ("20_wifi", "-57.0", "dBm", None, "mdi:wifi"),
     }
 
     for (sensor_id, data) in sensors.items():