From 398061706ca54c3f6c7567dc3b8459aa5e2534e3 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Wed, 20 Oct 2021 20:28:48 +0200
Subject: [PATCH] Correct unit_of_measurement for statistics sensor (#58023)

---
 homeassistant/components/statistics/sensor.py | 11 +++++------
 tests/components/statistics/test_sensor.py    | 13 +++++++++++++
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/statistics/sensor.py b/homeassistant/components/statistics/sensor.py
index 51f6478b4c0..6bf882c861d 100644
--- a/homeassistant/components/statistics/sensor.py
+++ b/homeassistant/components/statistics/sensor.py
@@ -150,10 +150,6 @@ class StatisticsSensor(SensorEntity):
             if (new_state := event.data.get("new_state")) is None:
                 return
 
-            self._unit_of_measurement = new_state.attributes.get(
-                ATTR_UNIT_OF_MEASUREMENT
-            )
-
             self._add_state_to_queue(new_state)
 
             self.async_schedule_update_ha_state(True)
@@ -171,7 +167,7 @@ class StatisticsSensor(SensorEntity):
 
             if "recorder" in self.hass.config.components:
                 # Only use the database if it's configured
-                self.hass.async_create_task(self._async_initialize_from_database())
+                self.hass.async_create_task(self._initialize_from_database())
 
         self.hass.bus.async_listen_once(
             EVENT_HOMEASSISTANT_START, async_stats_sensor_startup
@@ -195,6 +191,9 @@ class StatisticsSensor(SensorEntity):
                 self.entity_id,
                 new_state.state,
             )
+            return
+
+        self._unit_of_measurement = new_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
 
     @property
     def name(self):
@@ -355,7 +354,7 @@ class StatisticsSensor(SensorEntity):
                 self.hass, _scheduled_update, next_to_purge_timestamp
             )
 
-    async def _async_initialize_from_database(self):
+    async def _initialize_from_database(self):
         """Initialize the list of states from the database.
 
         The query will get the list of states in DESCENDING order so that we
diff --git a/tests/components/statistics/test_sensor.py b/tests/components/statistics/test_sensor.py
index bcbf13b8298..6cb53c9b93f 100644
--- a/tests/components/statistics/test_sensor.py
+++ b/tests/components/statistics/test_sensor.py
@@ -123,6 +123,18 @@ class TestStatisticsSensor(unittest.TestCase):
         assert self.change == state.attributes.get("change")
         assert self.average_change == state.attributes.get("average_change")
 
+        # Source sensor is unavailable, unit and state should not change
+        self.hass.states.set("sensor.test_monitored", "unavailable", {})
+        self.hass.block_till_done()
+        new_state = self.hass.states.get("sensor.test")
+        assert state == new_state
+
+        # Source sensor has a non float state, unit and state should not change
+        self.hass.states.set("sensor.test_monitored", "beer", {})
+        self.hass.block_till_done()
+        new_state = self.hass.states.get("sensor.test")
+        assert state == new_state
+
     def test_sampling_size(self):
         """Test rotation."""
         assert setup_component(
@@ -380,6 +392,7 @@ class TestStatisticsSensor(unittest.TestCase):
         # check if the result is as in test_sensor_source()
         state = self.hass.states.get("sensor.test")
         assert str(self.mean) == state.state
+        assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == TEMP_CELSIUS
 
     def test_initialize_from_database_with_maxage(self):
         """Test initializing the statistics from the database."""