From a0214d695aa79569e0bb0b4ead699e873ef1b94b Mon Sep 17 00:00:00 2001
From: Martin Hjelmare <marhje52@gmail.com>
Date: Fri, 25 Feb 2022 14:54:25 +0100
Subject: [PATCH] Clean up mysensors hass.data gateway access (#67233)

---
 homeassistant/components/mysensors/__init__.py |  9 +++------
 homeassistant/components/mysensors/const.py    |  2 --
 homeassistant/components/mysensors/gateway.py  | 11 -----------
 3 files changed, 3 insertions(+), 19 deletions(-)

diff --git a/homeassistant/components/mysensors/__init__.py b/homeassistant/components/mysensors/__init__.py
index 4d3c3046a89..dc4462b8d0e 100644
--- a/homeassistant/components/mysensors/__init__.py
+++ b/homeassistant/components/mysensors/__init__.py
@@ -43,7 +43,7 @@ from .const import (
     SensorType,
 )
 from .device import MySensorsDevice, get_mysensors_devices
-from .gateway import finish_setup, get_mysensors_gateway, gw_stop, setup_gateway
+from .gateway import finish_setup, gw_stop, setup_gateway
 from .helpers import on_unload
 
 _LOGGER = logging.getLogger(__name__)
@@ -244,7 +244,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Remove an instance of the MySensors integration."""
 
-    gateway = get_mysensors_gateway(hass, entry.entry_id)
+    gateway: BaseAsyncGateway = hass.data[DOMAIN][MYSENSORS_GATEWAYS][entry.entry_id]
 
     unload_ok = await hass.config_entries.async_unload_platforms(
         entry, PLATFORMS_WITH_ENTRY_SUPPORT
@@ -314,10 +314,7 @@ def setup_mysensors_platform(
             )
             continue
         gateway_id, node_id, child_id, value_type = dev_id
-        gateway: BaseAsyncGateway | None = get_mysensors_gateway(hass, gateway_id)
-        if not gateway:
-            _LOGGER.warning("Skipping setup of %s, no gateway found", dev_id)
-            continue
+        gateway: BaseAsyncGateway = hass.data[DOMAIN][MYSENSORS_GATEWAYS][gateway_id]
 
         if isinstance(device_class, dict):
             child = gateway.sensors[node_id].children[child_id]
diff --git a/homeassistant/components/mysensors/const.py b/homeassistant/components/mysensors/const.py
index ad25e601bbd..5f3cb6aed96 100644
--- a/homeassistant/components/mysensors/const.py
+++ b/homeassistant/components/mysensors/const.py
@@ -62,8 +62,6 @@ ValueType = str
 
 GatewayId = str
 # a unique id generated by config_flow.py and stored in the ConfigEntry as the entry id.
-#
-# Gateway may be fetched by giving the gateway id to get_mysensors_gateway()
 
 DevId = tuple[GatewayId, int, int, int]
 # describes the backend of a hass entity. Contents are: GatewayId, node_id, child_id, v_type as int
diff --git a/homeassistant/components/mysensors/gateway.py b/homeassistant/components/mysensors/gateway.py
index be0381ab74e..7f13035f55c 100644
--- a/homeassistant/components/mysensors/gateway.py
+++ b/homeassistant/components/mysensors/gateway.py
@@ -37,7 +37,6 @@ from .const import (
     CONF_VERSION,
     DOMAIN,
     MYSENSORS_GATEWAY_START_TASK,
-    MYSENSORS_GATEWAYS,
     ConfGatewayType,
     GatewayId,
 )
@@ -122,16 +121,6 @@ async def try_connect(
         return False
 
 
-def get_mysensors_gateway(
-    hass: HomeAssistant, gateway_id: GatewayId
-) -> BaseAsyncGateway | None:
-    """Return the Gateway for a given GatewayId."""
-    if MYSENSORS_GATEWAYS not in hass.data[DOMAIN]:
-        hass.data[DOMAIN][MYSENSORS_GATEWAYS] = {}
-    gateways = hass.data[DOMAIN].get(MYSENSORS_GATEWAYS)
-    return gateways.get(gateway_id)
-
-
 async def setup_gateway(
     hass: HomeAssistant, entry: ConfigEntry
 ) -> BaseAsyncGateway | None: