From 6644f62ad2c00f3d686f5042765a339f119a9e69 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 5 Sep 2022 13:56:27 -0500
Subject: [PATCH] Fix history stats device class when type is not time (#77855)

---
 .../components/history_stats/sensor.py        |  3 +-
 tests/components/history_stats/test_sensor.py | 45 ++++++++++++++++++-
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/history_stats/sensor.py b/homeassistant/components/history_stats/sensor.py
index a42c516f12b..642c327e29d 100644
--- a/homeassistant/components/history_stats/sensor.py
+++ b/homeassistant/components/history_stats/sensor.py
@@ -143,7 +143,6 @@ class HistoryStatsSensorBase(
 class HistoryStatsSensor(HistoryStatsSensorBase):
     """A HistoryStats sensor."""
 
-    _attr_device_class = SensorDeviceClass.DURATION
     _attr_state_class = SensorStateClass.MEASUREMENT
 
     def __init__(
@@ -157,6 +156,8 @@ class HistoryStatsSensor(HistoryStatsSensorBase):
         self._attr_native_unit_of_measurement = UNITS[sensor_type]
         self._type = sensor_type
         self._process_update()
+        if self._type == CONF_TYPE_TIME:
+            self._attr_device_class = SensorDeviceClass.DURATION
 
     @callback
     def _process_update(self) -> None:
diff --git a/tests/components/history_stats/test_sensor.py b/tests/components/history_stats/test_sensor.py
index 8907f381a6c..5de74f71d1e 100644
--- a/tests/components/history_stats/test_sensor.py
+++ b/tests/components/history_stats/test_sensor.py
@@ -9,7 +9,7 @@ import pytest
 
 from homeassistant import config as hass_config
 from homeassistant.components.history_stats import DOMAIN
-from homeassistant.const import SERVICE_RELOAD, STATE_UNKNOWN
+from homeassistant.const import ATTR_DEVICE_CLASS, SERVICE_RELOAD, STATE_UNKNOWN
 import homeassistant.core as ha
 from homeassistant.helpers.entity_component import async_update_entity
 from homeassistant.setup import async_setup_component, setup_component
@@ -1496,3 +1496,46 @@ async def test_end_time_with_microseconds_zeroed(time_zone, hass, recorder_mock)
         async_fire_time_changed(hass, rolled_to_next_day_plus_18)
         await hass.async_block_till_done()
         assert hass.states.get("sensor.heatpump_compressor_today").state == "16.0"
+
+
+async def test_device_classes(hass, recorder_mock):
+    """Test the device classes."""
+    await async_setup_component(
+        hass,
+        "sensor",
+        {
+            "sensor": [
+                {
+                    "platform": "history_stats",
+                    "entity_id": "binary_sensor.test_id",
+                    "name": "time",
+                    "state": "on",
+                    "start": "{{ as_timestamp(now()) - 3600 }}",
+                    "end": "{{ as_timestamp(now()) + 3600 }}",
+                    "type": "time",
+                },
+                {
+                    "platform": "history_stats",
+                    "entity_id": "binary_sensor.test_id",
+                    "name": "count",
+                    "state": "on",
+                    "start": "{{ as_timestamp(now()) - 3600 }}",
+                    "end": "{{ as_timestamp(now()) + 3600 }}",
+                    "type": "count",
+                },
+                {
+                    "platform": "history_stats",
+                    "entity_id": "binary_sensor.test_id",
+                    "name": "ratio",
+                    "state": "on",
+                    "start": "{{ as_timestamp(now()) - 3600 }}",
+                    "end": "{{ as_timestamp(now()) + 3600 }}",
+                    "type": "ratio",
+                },
+            ]
+        },
+    )
+    await hass.async_block_till_done()
+    assert hass.states.get("sensor.time").attributes[ATTR_DEVICE_CLASS] == "duration"
+    assert ATTR_DEVICE_CLASS not in hass.states.get("sensor.ratio").attributes
+    assert ATTR_DEVICE_CLASS not in hass.states.get("sensor.count").attributes