From d85d1a65a7a00670ce18fdff7483930e440334cf Mon Sep 17 00:00:00 2001
From: Martin Hjelmare <marhje52@gmail.com>
Date: Sun, 7 Mar 2021 14:20:21 +0100
Subject: [PATCH] Fix mysensors unload clean up (#47541)

---
 .../components/mysensors/__init__.py          | 21 +++---------------
 homeassistant/components/mysensors/gateway.py | 18 +++++++++++----
 homeassistant/components/mysensors/helpers.py | 22 ++++++++++++++++++-
 3 files changed, 38 insertions(+), 23 deletions(-)

diff --git a/homeassistant/components/mysensors/__init__.py b/homeassistant/components/mysensors/__init__.py
index bcab1ed86b4..7dd118099f7 100644
--- a/homeassistant/components/mysensors/__init__.py
+++ b/homeassistant/components/mysensors/__init__.py
@@ -38,11 +38,11 @@ from .const import (
     MYSENSORS_ON_UNLOAD,
     PLATFORMS_WITH_ENTRY_SUPPORT,
     DevId,
-    GatewayId,
     SensorType,
 )
 from .device import MySensorsDevice, MySensorsEntity, get_mysensors_devices
 from .gateway import finish_setup, get_mysensors_gateway, gw_stop, setup_gateway
+from .helpers import on_unload
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -253,29 +253,14 @@ async def async_unload_entry(hass: HomeAssistantType, entry: ConfigEntry) -> boo
         for fnct in hass.data[DOMAIN][key]:
             fnct()
 
+        hass.data[DOMAIN].pop(key)
+
     del hass.data[DOMAIN][MYSENSORS_GATEWAYS][entry.entry_id]
 
     await gw_stop(hass, entry, gateway)
     return True
 
 
-async def on_unload(
-    hass: HomeAssistantType, entry: Union[ConfigEntry, GatewayId], fnct: Callable
-) -> None:
-    """Register a callback to be called when entry is unloaded.
-
-    This function is used by platforms to cleanup after themselves
-    """
-    if isinstance(entry, GatewayId):
-        uniqueid = entry
-    else:
-        uniqueid = entry.entry_id
-    key = MYSENSORS_ON_UNLOAD.format(uniqueid)
-    if key not in hass.data[DOMAIN]:
-        hass.data[DOMAIN][key] = []
-    hass.data[DOMAIN][key].append(fnct)
-
-
 @callback
 def setup_mysensors_platform(
     hass: HomeAssistant,
diff --git a/homeassistant/components/mysensors/gateway.py b/homeassistant/components/mysensors/gateway.py
index b6797cafb37..a7f3a053d3f 100644
--- a/homeassistant/components/mysensors/gateway.py
+++ b/homeassistant/components/mysensors/gateway.py
@@ -31,7 +31,12 @@ from .const import (
     GatewayId,
 )
 from .handler import HANDLERS
-from .helpers import discover_mysensors_platform, validate_child, validate_node
+from .helpers import (
+    discover_mysensors_platform,
+    on_unload,
+    validate_child,
+    validate_node,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -260,8 +265,8 @@ async def _discover_persistent_devices(
 
 async def gw_stop(hass, entry: ConfigEntry, gateway: BaseAsyncGateway):
     """Stop the gateway."""
-    connect_task = hass.data[DOMAIN].get(
-        MYSENSORS_GATEWAY_START_TASK.format(entry.entry_id)
+    connect_task = hass.data[DOMAIN].pop(
+        MYSENSORS_GATEWAY_START_TASK.format(entry.entry_id), None
     )
     if connect_task is not None and not connect_task.done():
         connect_task.cancel()
@@ -288,7 +293,12 @@ async def _gw_start(
     async def stop_this_gw(_: Event):
         await gw_stop(hass, entry, gateway)
 
-    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, stop_this_gw)
+    await on_unload(
+        hass,
+        entry.entry_id,
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, stop_this_gw),
+    )
+
     if entry.data[CONF_DEVICE] == MQTT_COMPONENT:
         # Gatways connected via mqtt doesn't send gateway ready message.
         return
diff --git a/homeassistant/components/mysensors/helpers.py b/homeassistant/components/mysensors/helpers.py
index 4452dd0575b..0b8dc361158 100644
--- a/homeassistant/components/mysensors/helpers.py
+++ b/homeassistant/components/mysensors/helpers.py
@@ -2,16 +2,18 @@
 from collections import defaultdict
 from enum import IntEnum
 import logging
-from typing import DefaultDict, Dict, List, Optional, Set
+from typing import Callable, DefaultDict, Dict, List, Optional, Set, Union
 
 from mysensors import BaseAsyncGateway, Message
 from mysensors.sensor import ChildSensor
 import voluptuous as vol
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.typing import HomeAssistantType
 from homeassistant.util.decorator import Registry
 
 from .const import (
@@ -20,6 +22,7 @@ from .const import (
     DOMAIN,
     FLAT_PLATFORM_TYPES,
     MYSENSORS_DISCOVERY,
+    MYSENSORS_ON_UNLOAD,
     TYPE_TO_PLATFORMS,
     DevId,
     GatewayId,
@@ -31,6 +34,23 @@ _LOGGER = logging.getLogger(__name__)
 SCHEMAS = Registry()
 
 
+async def on_unload(
+    hass: HomeAssistantType, entry: Union[ConfigEntry, GatewayId], fnct: Callable
+) -> None:
+    """Register a callback to be called when entry is unloaded.
+
+    This function is used by platforms to cleanup after themselves.
+    """
+    if isinstance(entry, GatewayId):
+        uniqueid = entry
+    else:
+        uniqueid = entry.entry_id
+    key = MYSENSORS_ON_UNLOAD.format(uniqueid)
+    if key not in hass.data[DOMAIN]:
+        hass.data[DOMAIN][key] = []
+    hass.data[DOMAIN][key].append(fnct)
+
+
 @callback
 def discover_mysensors_platform(
     hass: HomeAssistant, gateway_id: GatewayId, platform: str, new_devices: List[DevId]