From ade6fd02157b12bee70992bbabc25f6a72ebb579 Mon Sep 17 00:00:00 2001
From: Whitney Young <wbyoung@users.noreply.github.com>
Date: Thu, 8 Aug 2024 15:20:40 -0700
Subject: [PATCH] Get the proxy as well as the gateway

---
 custom_components/zha_toolkit/__init__.py | 10 +++++-----
 custom_components/zha_toolkit/utils.py    | 21 ++++++++++++++++++++-
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/custom_components/zha_toolkit/__init__.py b/custom_components/zha_toolkit/__init__.py
index 5a4057d..6a59ac4 100644
--- a/custom_components/zha_toolkit/__init__.py
+++ b/custom_components/zha_toolkit/__init__.py
@@ -1,6 +1,6 @@
 import importlib
 import logging
-from typing import Optional
+from typing import Any, Optional
 
 import homeassistant.helpers.config_validation as cv
 import voluptuous as vol
@@ -677,8 +677,8 @@ async def toolkit_service(service):
         global LOADED_VERSION  # pylint: disable=global-variable-not-assigned
 
         zha = hass_ref.data["zha"]
-        zha_gw: Optional[ZHAGateway] = None
-        zha_gw = u.get_zha_gateway(hass)
+        zha_gw: Optional[ZHAGateway] = u.get_zha_gateway(hass)
+        zha_gw_hass: Any = u.get_zha_gateway_hass(hass)
 
         if zha_gw is None:
             LOGGER.error(
@@ -725,7 +725,7 @@ async def toolkit_service(service):
 
         app = zha_gw.application_controller  # type: ignore
 
-        ieee = await u.get_ieee(app, zha_gw, ieee_str)
+        ieee = await u.get_ieee(app, zha_gw_hass, ieee_str)
 
         slickParams = params.copy()
         for k in params:
@@ -783,7 +783,7 @@ async def toolkit_service(service):
         try:
             handler_result = await handler(
                 zha_gw.application_controller,  # type: ignore
-                zha_gw,
+                zha_gw_hass,
                 ieee,
                 cmd,
                 cmd_data,
diff --git a/custom_components/zha_toolkit/utils.py b/custom_components/zha_toolkit/utils.py
index 3a471db..5df20ea 100644
--- a/custom_components/zha_toolkit/utils.py
+++ b/custom_components/zha_toolkit/utils.py
@@ -19,7 +19,12 @@
     from homeassistant.components.zha.core.gateway import ZHAGateway
 
 from homeassistant.components import zha
-from homeassistant.components.zha import helpers as zha_helpers
+
+try:
+    from homeassistant.components.zha import helpers as zha_helpers
+except ImportError:
+    zha_helpers = None
+
 from homeassistant.util import dt as dt_util
 from pkg_resources import get_distribution, parse_version
 from zigpy import types as t
@@ -62,6 +67,20 @@ def get_zha_gateway(hass: HomeAssistant) -> ZHAGateway:
     return zha.gateway
 
 
+def get_zha_gateway_hass(
+    hass: HomeAssistant,
+) -> ZHAGateway | zha_helpers.ZHAGatewayProxy:
+    """
+    Get the ZHA gateway proxy object.
+
+    Fallback to the gateway object prior to 2024.8 which still has an attached
+    HASS object.
+    """
+    if parse_version(HA_VERSION) >= parse_version("2024.8"):
+        return zha_helpers.get_zha_gateway_proxy(hass)
+    return get_zha_gateway(hass)
+
+
 def getHaVersion() -> str:
     """Get HA Version"""
     return HA_VERSION