From c1ddde376480b7353503a48c5765f4829a6e9d25 Mon Sep 17 00:00:00 2001
From: RoboMagus <68224306+RoboMagus@users.noreply.github.com>
Date: Wed, 25 May 2022 08:44:08 +0200
Subject: [PATCH] Check if attributes are present in new_state before accessing
 them (#71967)

* Check if attributes are present in new_state before accessing them.

* Early return if new state is None|Unknown|Unavailable

* Removed whitespace at line endings. +black run

* Update test for coverage
---
 homeassistant/components/integration/sensor.py | 15 +++++++++------
 tests/components/integration/test_sensor.py    | 10 ++++++++++
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py
index 8cf5da5bc7c..5c982c6ec5e 100644
--- a/homeassistant/components/integration/sensor.py
+++ b/homeassistant/components/integration/sensor.py
@@ -191,11 +191,16 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
             old_state = event.data.get("old_state")
             new_state = event.data.get("new_state")
 
+            if new_state is None or new_state.state in (
+                STATE_UNKNOWN,
+                STATE_UNAVAILABLE,
+            ):
+                return
+
             # We may want to update our state before an early return,
             # based on the source sensor's unit_of_measurement
             # or device_class.
             update_state = False
-
             unit = new_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
             if unit is not None:
                 new_unit_of_measurement = self._unit_template.format(unit)
@@ -214,11 +219,9 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
             if update_state:
                 self.async_write_ha_state()
 
-            if (
-                old_state is None
-                or new_state is None
-                or old_state.state in (STATE_UNKNOWN, STATE_UNAVAILABLE)
-                or new_state.state in (STATE_UNKNOWN, STATE_UNAVAILABLE)
+            if old_state is None or old_state.state in (
+                STATE_UNKNOWN,
+                STATE_UNAVAILABLE,
             ):
                 return
 
diff --git a/tests/components/integration/test_sensor.py b/tests/components/integration/test_sensor.py
index 90c9e6f0fbc..e4173d62eb4 100644
--- a/tests/components/integration/test_sensor.py
+++ b/tests/components/integration/test_sensor.py
@@ -9,6 +9,7 @@ from homeassistant.const import (
     ENERGY_WATT_HOUR,
     POWER_KILO_WATT,
     POWER_WATT,
+    STATE_UNAVAILABLE,
     STATE_UNKNOWN,
     TIME_SECONDS,
 )
@@ -350,6 +351,15 @@ async def test_units(hass):
     # they became valid
     assert state.attributes.get("unit_of_measurement") == ENERGY_WATT_HOUR
 
+    # When source state goes to None / Unknown, expect an early exit without
+    # changes to the state or unit_of_measurement
+    hass.states.async_set(entity_id, STATE_UNAVAILABLE, None)
+    await hass.async_block_till_done()
+
+    new_state = hass.states.get("sensor.integration")
+    assert state == new_state
+    assert state.attributes.get("unit_of_measurement") == ENERGY_WATT_HOUR
+
 
 async def test_device_class(hass):
     """Test integration sensor units using a power source."""