diff --git a/config_flow.py b/config_flow.py
index f1bb831..d350ca7 100644
--- a/config_flow.py
+++ b/config_flow.py
@@ -19,15 +19,15 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_get_zha_devices(hass):
     """Get ZHA devices."""
-    if "zha" not in hass.data:
+    if not hasattr(hass.data.get("zha"), "core"):
         return []
 
-    zha_gateway = hass.data["zha"].get("zha_gateway")
+    zha_gateway = hass.data["zha"].core.gateway
     if not zha_gateway:
         return []
 
     return [
-        (str(device.ieee), f"{device.name} ({device.ieee})")
+        (str(device.ieee), f"{device.name} ({device.ieee})"
         for device in zha_gateway.devices.values()
     ]
 
diff --git a/init.py b/init.py
index c28f28d..2629f09 100644
--- a/init.py
+++ b/init.py
@@ -7,7 +7,6 @@ from typing import Any
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_IEEE
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import (
@@ -68,10 +67,10 @@ class LinkyTariffCoordinator(DataUpdateCoordinator):
     async def _async_update_data(self) -> dict[str, Any]:
         """Fetch data from ZHA device."""
         try:
-            if "zha" not in self.hass.data:
+            if not hasattr(self.hass.data.get("zha"), "core"):
                 raise ValueError("ZHA integration not loaded")
 
-            zha_gateway = self.hass.data["zha"].get("zha_gateway")
+            zha_gateway = self.hass.data["zha"].core.gateway
             if not zha_gateway:
                 raise ValueError("ZHA gateway not available")