From c5f7d7ae857d45f1e71fc97557a7a6e5225bf94c Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Fri, 30 Dec 2022 19:08:11 +0100
Subject: [PATCH] Only reflect unavailable state in DSMR when disconnected
 (#84862)

* Only reflect unavailable state in DSMR when disonnected

* Addressreview comment
---
 homeassistant/components/dsmr/sensor.py | 34 +++++++++++++++++++------
 tests/components/dsmr/test_sensor.py    |  5 ++++
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/homeassistant/components/dsmr/sensor.py b/homeassistant/components/dsmr/sensor.py
index 7147655549b4d5..36f734273f2f73 100644
--- a/homeassistant/components/dsmr/sensor.py
+++ b/homeassistant/components/dsmr/sensor.py
@@ -401,7 +401,7 @@ async def async_setup_entry(
     )
 
     @Throttle(min_time_between_updates)
-    def update_entities_telegram(telegram: dict[str, DSMRObject]) -> None:
+    def update_entities_telegram(telegram: dict[str, DSMRObject] | None) -> None:
         """Update entities with latest telegram and trigger state update."""
         # Make all device entities aware of new telegram
         for entity in entities:
@@ -445,6 +445,11 @@ async def connect_and_reconnect() -> None:
 
         while hass.state == CoreState.not_running or hass.is_running:
             # Start DSMR asyncio.Protocol reader
+
+            # Reflect connected state in devices state by setting an
+            # empty telegram resulting in `unknown` states
+            update_entities_telegram({})
+
             try:
                 transport, protocol = await hass.loop.create_task(reader_factory())
 
@@ -472,8 +477,8 @@ def close_transport(_event: EventType) -> None:
                 protocol = None
 
                 # Reflect disconnect state in devices state by setting an
-                # empty telegram resulting in `unknown` states
-                update_entities_telegram({})
+                # None telegram resulting in `unavailable` states
+                update_entities_telegram(None)
 
                 # throttle reconnect attempts
                 await asyncio.sleep(
@@ -487,11 +492,19 @@ def close_transport(_event: EventType) -> None:
                 transport = None
                 protocol = None
 
+                # Reflect disconnect state in devices state by setting an
+                # None telegram resulting in `unavailable` states
+                update_entities_telegram(None)
+
                 # throttle reconnect attempts
                 await asyncio.sleep(
                     entry.data.get(CONF_RECONNECT_INTERVAL, DEFAULT_RECONNECT_INTERVAL)
                 )
             except CancelledError:
+                # Reflect disconnect state in devices state by setting an
+                # None telegram resulting in `unavailable` states
+                update_entities_telegram(None)
+
                 if stop_listener and (
                     hass.state == CoreState.not_running or hass.is_running
                 ):
@@ -534,7 +547,7 @@ def __init__(
         """Initialize entity."""
         self.entity_description = entity_description
         self._entry = entry
-        self.telegram: dict[str, DSMRObject] = {}
+        self.telegram: dict[str, DSMRObject] | None = {}
 
         device_serial = entry.data[CONF_SERIAL_ID]
         device_name = DEVICE_NAME_ELECTRICITY
@@ -551,16 +564,21 @@ def __init__(
         self._attr_unique_id = f"{device_serial}_{entity_description.key}"
 
     @callback
-    def update_data(self, telegram: dict[str, DSMRObject]) -> None:
+    def update_data(self, telegram: dict[str, DSMRObject] | None) -> None:
         """Update data."""
         self.telegram = telegram
-        if self.hass and self.entity_description.obis_reference in self.telegram:
+        if self.hass and (
+            telegram is None or self.entity_description.obis_reference in telegram
+        ):
             self.async_write_ha_state()
 
     def get_dsmr_object_attr(self, attribute: str) -> str | None:
         """Read attribute from last received telegram for this DSMR object."""
         # Make sure telegram contains an object for this entities obis
-        if self.entity_description.obis_reference not in self.telegram:
+        if (
+            self.telegram is None
+            or self.entity_description.obis_reference not in self.telegram
+        ):
             return None
 
         # Get the attribute value if the object has it
@@ -571,7 +589,7 @@ def get_dsmr_object_attr(self, attribute: str) -> str | None:
     @property
     def available(self) -> bool:
         """Entity is only available if there is a telegram."""
-        return bool(self.telegram)
+        return self.telegram is not None
 
     @property
     def native_value(self) -> StateType:
diff --git a/tests/components/dsmr/test_sensor.py b/tests/components/dsmr/test_sensor.py
index 31d3496e3a7b74..ee0ffa5db5fef6 100644
--- a/tests/components/dsmr/test_sensor.py
+++ b/tests/components/dsmr/test_sensor.py
@@ -24,6 +24,7 @@
     ATTR_UNIT_OF_MEASUREMENT,
     ENERGY_KILO_WATT_HOUR,
     STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
     VOLUME_CUBIC_METERS,
     UnitOfPower,
 )
@@ -783,6 +784,10 @@ async def wait_closed():
 
     assert connection_factory.call_count == 1
 
+    state = hass.states.get("sensor.electricity_meter_power_consumption")
+    assert state
+    assert state.state == STATE_UNKNOWN
+
     # indicate disconnect, release wait lock and allow reconnect to happen
     closed.set()
     # wait for lock set to resolve