Update ZHA state handling (#21866)

* make device available if it was seen within 2 hours

* more state restore

* cleanup init

* clean up storage stuff

* fix tests

* update state handling

homeassistant/components/zha/__init__.py

@@ -23,10 +23,11 @@ from .core.const import (
     COMPONENTS, CONF_BAUDRATE, CONF_DATABASE, CONF_DEVICE_CONFIG,
     CONF_RADIO_TYPE, CONF_USB_PATH, DATA_ZHA, DATA_ZHA_BRIDGE_ID,
     DATA_ZHA_CONFIG, DATA_ZHA_CORE_COMPONENT, DATA_ZHA_DISPATCHERS,
-    DATA_ZHA_RADIO, DEFAULT_BAUDRATE, DEFAULT_DATABASE_NAME,
+    DATA_ZHA_RADIO, DEFAULT_BAUDRATE, DEFAULT_DATABASE_NAME, DATA_ZHA_GATEWAY,
     DEFAULT_RADIO_TYPE, DOMAIN, RadioType, DATA_ZHA_CORE_EVENTS, ENABLE_QUIRKS)
 from .core.gateway import establish_device_mappings
 from .core.channels.registry import populate_channel_registry
+from .core.store import async_get_registry
 
 REQUIREMENTS = [
     'bellows-homeassistant==0.7.1',
@@ -146,7 +147,8 @@ async def async_setup_entry(hass, config_entry):
         ClusterPersistingListener
     )
 
-    zha_gateway = ZHAGateway(hass, config)
+    zha_storage = await async_get_registry(hass)
+    zha_gateway = ZHAGateway(hass, config, zha_storage)
 
     # Patch handle_message until zigpy can provide an event here
     def handle_message(sender, is_reply, profile, cluster,
@@ -192,11 +194,14 @@ async def async_setup_entry(hass, config_entry):
 
     api.async_load_api(hass, application_controller, zha_gateway)
 
-    def zha_shutdown(event):
+    async def async_zha_shutdown(event):
-        """Close radio."""
+        """Handle shutdown tasks."""
+        await hass.data[DATA_ZHA][
+            DATA_ZHA_GATEWAY].async_update_device_storage()
         hass.data[DATA_ZHA][DATA_ZHA_RADIO].close()
 
-    hass.bus.async_listen_once(ha_const.EVENT_HOMEASSISTANT_STOP, zha_shutdown)
+    hass.bus.async_listen_once(
+        ha_const.EVENT_HOMEASSISTANT_STOP, async_zha_shutdown)
     return True
 
 

homeassistant/components/zha/binary_sensor.py

@@ -7,6 +7,8 @@ at https://home-assistant.io/components/binary_sensor.zha/
 import logging
 
 from homeassistant.components.binary_sensor import DOMAIN, BinarySensorDevice
+from homeassistant.const import STATE_ON
+from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from .core.const import (
     DATA_ZHA, DATA_ZHA_DISPATCHERS, ZHA_DISCOVERY_NEW, ON_OFF_CHANNEL,
@@ -126,6 +128,14 @@ class BinarySensor(ZhaEntity, BinarySensorDevice):
             await self.async_accept_signal(
                 self._attr_channel, SIGNAL_ATTR_UPDATED, self.async_set_state)
 
+    @callback
+    def async_restore_last_state(self, last_state):
+        """Restore previous state."""
+        super().async_restore_last_state(last_state)
+        self._state = last_state.state == STATE_ON
+        if 'level' in last_state.attributes:
+            self._level = last_state.attributes['level']
+
     @property
     def is_on(self) -> bool:
         """Return if the switch is on based on the statemachine."""
@@ -166,3 +176,21 @@ class BinarySensor(ZhaEntity, BinarySensorDevice):
                 ATTR_LEVEL: self._state and self._level or 0
             })
         return self._device_state_attributes
+
+    async def async_update(self):
+        """Attempt to retrieve on off state from the binary sensor."""
+        await super().async_update()
+        if self._level_channel:
+            self._level = await self._level_channel.get_attribute_value(
+                'current_level')
+        if self._on_off_channel:
+            self._state = await self._on_off_channel.get_attribute_value(
+                'on_off')
+        if self._zone_channel:
+            value = await self._zone_channel.get_attribute_value(
+                'zone_status')
+            if value is not None:
+                self._state = value & 3
+        if self._attr_channel:
+            self._state = await self._attr_channel.get_attribute_value(
+                self._attr_channel.value_attribute)

homeassistant/components/zha/core/channels/__init__.py

@@ -20,7 +20,6 @@ from ..const import (
     CLUSTER_REPORT_CONFIGS, REPORT_CONFIG_DEFAULT, SIGNAL_ATTR_UPDATED,
     ATTRIBUTE_CHANNEL, EVENT_RELAY_CHANNEL, ZDO_CHANNEL
 )
-from ..store import async_get_registry
 
 NODE_DESCRIPTOR_REQUEST = 0x0002
 MAINS_POWERED = 1
@@ -221,14 +220,14 @@ class AttributeListeningChannel(ZigbeeChannel):
         self.name = ATTRIBUTE_CHANNEL
         attr = self._report_config[0].get('attr')
         if isinstance(attr, str):
-            self._value_attribute = get_attr_id_by_name(self.cluster, attr)
+            self.value_attribute = get_attr_id_by_name(self.cluster, attr)
         else:
-            self._value_attribute = attr
+            self.value_attribute = attr
 
     @callback
     def attribute_updated(self, attrid, value):
         """Handle attribute updates on this cluster."""
-        if attrid == self._value_attribute:
+        if attrid == self.value_attribute:
             async_dispatcher_send(
                 self._zha_device.hass,
                 "{}_{}".format(self.unique_id, SIGNAL_ATTR_UPDATED),
@@ -288,8 +287,8 @@ class ZDOChannel:
 
     async def async_initialize(self, from_cache):
         """Initialize channel."""
-        entry = (await async_get_registry(
+        entry = self._zha_device.gateway.zha_storage.async_get_or_create(
-            self._zha_device.hass)).async_get_or_create(self._zha_device)
+            self._zha_device)
         _LOGGER.debug("entry loaded from storage: %s", entry)
         if entry is not None:
             self.power_source = entry.power_source
@@ -303,8 +302,8 @@ class ZDOChannel:
             # this previously so lets set it up so users don't have
             # to reconfigure every device.
             await self.async_get_node_descriptor(False)
-            entry = (await async_get_registry(
+            entry = self._zha_device.gateway.zha_storage.async_update(
-                self._zha_device.hass)).async_update(self._zha_device)
+                self._zha_device)
             _LOGGER.debug("entry after getting node desc in init: %s", entry)
         self._status = ChannelStatus.INITIALIZED
 

homeassistant/components/zha/core/device.py

@@ -20,7 +20,6 @@ from .const import (
     QUIRK_CLASS, ZDO_CHANNEL, MANUFACTURER_CODE, POWER_SOURCE
 )
 from .channels import EventRelayChannel, ZDOChannel
-from .store import async_get_registry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -69,6 +68,7 @@ class ZHADevice:
             self._zigpy_device.__class__.__module__,
             self._zigpy_device.__class__.__name__
         )
+        self._power_source = None
         self.status = DeviceStatus.CREATED
 
     @property
@@ -120,7 +120,9 @@ class ZHADevice:
 
     @property
     def power_source(self):
-        """Return True if sensor is available."""
+        """Return the power source for the device."""
+        if self._power_source is not None:
+            return self._power_source
         if ZDO_CHANNEL in self.cluster_channels:
             return self.cluster_channels.get(ZDO_CHANNEL).power_source
         return None
@@ -145,6 +147,14 @@ class ZHADevice:
         """Return True if sensor is available."""
         return self._available
 
+    def set_available(self, available):
+        """Set availability from restore and prevent signals."""
+        self._available = available
+
+    def set_power_source(self, power_source):
+        """Set the power source."""
+        self._power_source = power_source
+
     def update_available(self, available):
         """Set sensor availability."""
         if self._available != available and available:
@@ -195,8 +205,7 @@ class ZHADevice:
         _LOGGER.debug('%s: started configuration', self.name)
         await self._execute_channel_tasks('async_configure')
         _LOGGER.debug('%s: completed configuration', self.name)
-        entry = (await async_get_registry(
+        entry = self.gateway.zha_storage.async_create_or_update(self)
-            self.hass)).async_create_or_update(self)
         _LOGGER.debug('%s: stored in registry: %s', self.name, entry)
 
     async def async_initialize(self, from_cache=False):
@@ -253,6 +262,11 @@ class ZHADevice:
         if self._unsub:
             self._unsub()
 
+    @callback
+    def async_update_last_seen(self, last_seen):
+        """Set last seen on the zigpy device."""
+        self._zigpy_device.last_seen = last_seen
+
     @callback
     def async_get_clusters(self):
         """Get all clusters for this device."""

homeassistant/components/zha/core/gateway.py

@@ -45,13 +45,14 @@ EntityReference = collections.namedtuple(
 class ZHAGateway:
     """Gateway that handles events that happen on the ZHA Zigbee network."""
 
-    def __init__(self, hass, config):
+    def __init__(self, hass, config, zha_storage):
         """Initialize the gateway."""
         self._hass = hass
         self._config = config
         self._component = EntityComponent(_LOGGER, DOMAIN, hass)
         self._devices = {}
         self._device_registry = collections.defaultdict(list)
+        self.zha_storage = zha_storage
         hass.data[DATA_ZHA][DATA_ZHA_CORE_COMPONENT] = self._component
         hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] = self
 
@@ -125,12 +126,16 @@ class ZHAGateway:
         )
 
     @callback
-    def _async_get_or_create_device(self, zigpy_device):
+    def _async_get_or_create_device(self, zigpy_device, is_new_join):
         """Get or create a ZHA device."""
         zha_device = self._devices.get(zigpy_device.ieee)
         if zha_device is None:
             zha_device = ZHADevice(self._hass, zigpy_device, self)
             self._devices[zigpy_device.ieee] = zha_device
+        if not is_new_join:
+            entry = self.zha_storage.async_get_or_create(zha_device)
+            zha_device.async_update_last_seen(entry.last_seen)
+            zha_device.set_power_source(entry.power_source)
         return zha_device
 
     @callback
@@ -149,9 +154,16 @@ class ZHAGateway:
             if device.status is DeviceStatus.INITIALIZED:
                 device.update_available(True)
 
+    async def async_update_device_storage(self):
+        """Update the devices in the store."""
+        for device in self.devices.values():
+            self.zha_storage.async_update(device)
+        await self.zha_storage.async_save()
+
     async def async_device_initialized(self, device, is_new_join):
         """Handle device joined and basic information discovered (async)."""
-        zha_device = self._async_get_or_create_device(device)
+        zha_device = self._async_get_or_create_device(device, is_new_join)
+
         discovery_infos = []
         for endpoint_id, endpoint in device.endpoints.items():
             self._async_process_endpoint(
@@ -162,10 +174,11 @@ class ZHAGateway:
         if is_new_join:
             # configure the device
             await zha_device.async_configure()
-        elif not zha_device.available and zha_device.power_source is not None\
+            zha_device.update_available(True)
+        elif zha_device.power_source is not None\
                 and zha_device.power_source == MAINS_POWERED:
-            # the device is currently marked unavailable and it isn't a battery
+            # the device isn't a battery powered device so we should be able
-            # powered device so we should be able to update it now
+            # to update it now
             _LOGGER.debug(
                 "attempting to request fresh state for %s %s",
                 zha_device.name,
@@ -187,11 +200,6 @@ class ZHAGateway:
         device_entity = _async_create_device_entity(zha_device)
         await self._component.async_add_entities([device_entity])
 
-        if is_new_join:
-            # because it's a new join we can immediately mark the device as
-            # available. We do it here because the entities didn't exist above
-            zha_device.update_available(True)
-
     @callback
     def _async_process_endpoint(
             self, endpoint_id, endpoint, discovery_infos, device, zha_device,

homeassistant/components/zha/core/store.py

@@ -28,6 +28,7 @@ class ZhaDeviceEntry:
     ieee = attr.ib(type=str, default=None)
     power_source = attr.ib(type=int, default=None)
     manufacturer_code = attr.ib(type=int, default=None)
+    last_seen = attr.ib(type=float, default=None)
 
 
 class ZhaDeviceStorage:
@@ -46,7 +47,8 @@ class ZhaDeviceStorage:
             name=device.name,
             ieee=str(device.ieee),
             power_source=device.power_source,
-            manufacturer_code=device.manufacturer_code
+            manufacturer_code=device.manufacturer_code,
+            last_seen=device.last_seen
 
         )
         self.devices[device_entry.ieee] = device_entry
@@ -68,10 +70,13 @@ class ZhaDeviceStorage:
             return self.async_update(device)
         return self.async_create(device)
 
-    async def async_delete(self, ieee: str) -> None:
+    @callback
+    def async_delete(self, device) -> None:
         """Delete ZhaDeviceEntry."""
-        del self.devices[ieee]
+        ieee_str = str(device.ieee)
-        self.async_schedule_save()
+        if ieee_str in self.devices:
+            del self.devices[ieee_str]
+            self.async_schedule_save()
 
     @callback
     def async_update(self, device) -> ZhaDeviceEntry:
@@ -87,6 +92,8 @@ class ZhaDeviceStorage:
         if device.manufacturer_code != old.manufacturer_code:
             changes['manufacturer_code'] = device.manufacturer_code
 
+        changes['last_seen'] = device.last_seen
+
         new = self.devices[ieee_str] = attr.evolve(old, **changes)
         self.async_schedule_save()
         return new
@@ -103,7 +110,9 @@ class ZhaDeviceStorage:
                     name=device['name'],
                     ieee=device['ieee'],
                     power_source=device['power_source'],
-                    manufacturer_code=device['manufacturer_code']
+                    manufacturer_code=device['manufacturer_code'],
+                    last_seen=device['last_seen'] if 'last_seen' in device
+                    else None
                 )
 
         self.devices = devices
@@ -113,6 +122,10 @@ class ZhaDeviceStorage:
         """Schedule saving the registry of zha devices."""
         self._store.async_delay_save(self._data_to_save, SAVE_DELAY)
 
+    async def async_save(self) -> None:
+        """Save the registry of zha devices."""
+        await self._store.async_save(self._data_to_save())
+
     @callback
     def _data_to_save(self) -> dict:
         """Return data for the registry of zha devices to store in a file."""
@@ -124,6 +137,7 @@ class ZhaDeviceStorage:
                 'ieee': entry.ieee,
                 'power_source': entry.power_source,
                 'manufacturer_code': entry.manufacturer_code,
+                'last_seen': entry.last_seen
             } for entry in self.devices.values()
         ]
 

homeassistant/components/zha/device_entity.py

@@ -98,6 +98,7 @@ class ZhaDeviceEntity(ZhaEntity):
     async def async_added_to_hass(self):
         """Run when about to be added to hass."""
         await super().async_added_to_hass()
+        await self.async_check_recently_seen()
         if self._battery_channel:
             await self.async_accept_signal(
                 self._battery_channel, SIGNAL_STATE_ATTR,

homeassistant/components/zha/entity.py

@@ -6,23 +6,28 @@ https://home-assistant.io/components/zha/
 """
 
 import logging
+import time
 
+from homeassistant.core import callback
 from homeassistant.helpers import entity
 from homeassistant.helpers.device_registry import CONNECTION_ZIGBEE
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.util import slugify
 
 from .core.const import (
     DOMAIN, ATTR_MANUFACTURER, DATA_ZHA, DATA_ZHA_BRIDGE_ID, MODEL, NAME,
     SIGNAL_REMOVE
 )
+from .core.channels import MAINS_POWERED
 
 _LOGGER = logging.getLogger(__name__)
 
 ENTITY_SUFFIX = 'entity_suffix'
+RESTART_GRACE_PERIOD = 7200  # 2 hours
 
 
-class ZhaEntity(entity.Entity):
+class ZhaEntity(RestoreEntity, entity.Entity):
     """A base class for ZHA entities."""
 
     _domain = None  # Must be overridden by subclasses
@@ -136,6 +141,7 @@ class ZhaEntity(entity.Entity):
     async def async_added_to_hass(self):
         """Run when about to be added to hass."""
         await super().async_added_to_hass()
+        await self.async_check_recently_seen()
         await self.async_accept_signal(
             None, "{}_{}".format(self.zha_device.available_signal, 'entity'),
             self.async_set_available,
@@ -149,11 +155,28 @@ class ZhaEntity(entity.Entity):
             self._zha_device.ieee, self.entity_id, self._zha_device,
             self.cluster_channels, self.device_info)
 
+    async def async_check_recently_seen(self):
+        """Check if the device was seen within the last 2 hours."""
+        last_state = await self.async_get_last_state()
+        if last_state and self._zha_device.last_seen and (
+                time.time() - self._zha_device.last_seen <
+                RESTART_GRACE_PERIOD):
+            self.async_set_available(True)
+            if self.zha_device.power_source != MAINS_POWERED:
+                # mains powered devices will get real time state
+                self.async_restore_last_state(last_state)
+            self._zha_device.set_available(True)
+
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect entity object when removed."""
         for unsub in self._unsubs:
             unsub()
 
+    @callback
+    def async_restore_last_state(self, last_state):
+        """Restore previous state."""
+        pass
+
     async def async_update(self):
         """Retrieve latest state."""
         for channel in self.cluster_channels:

homeassistant/components/zha/fan.py

@@ -6,6 +6,7 @@ at https://home-assistant.io/components/fan.zha/
 """
 import logging
 
+from homeassistant.core import callback
 from homeassistant.components.fan import (
     DOMAIN, SPEED_HIGH, SPEED_LOW, SPEED_MEDIUM, SPEED_OFF, SUPPORT_SET_SPEED,
     FanEntity)
@@ -92,6 +93,11 @@ class ZhaFan(ZhaEntity, FanEntity):
         await self.async_accept_signal(
             self._fan_channel, SIGNAL_ATTR_UPDATED, self.async_set_state)
 
+    @callback
+    def async_restore_last_state(self, last_state):
+        """Restore previous state."""
+        self._state = VALUE_TO_SPEED.get(last_state.state, last_state.state)
+
     @property
     def supported_features(self) -> int:
         """Flag supported features."""
@@ -139,3 +145,11 @@ class ZhaFan(ZhaEntity, FanEntity):
         """Set the speed of the fan."""
         await self._fan_channel.async_set_speed(SPEED_TO_VALUE[speed])
         self.async_set_state(speed)
+
+    async def async_update(self):
+        """Attempt to retrieve on off state from the fan."""
+        await super().async_update()
+        if self._fan_channel:
+            state = await self._fan_channel.get_attribute_value('fan_mode')
+            if state is not None:
+                self._state = VALUE_TO_SPEED.get(state, self._state)

homeassistant/components/zha/light.py

@@ -8,6 +8,8 @@ from datetime import timedelta
 import logging
 
 from homeassistant.components import light
+from homeassistant.const import STATE_ON
+from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 import homeassistant.util.color as color_util
 from .const import (
@@ -156,6 +158,17 @@ class Light(ZhaEntity, light.Light):
             await self.async_accept_signal(
                 self._level_channel, SIGNAL_SET_LEVEL, self.set_level)
 
+    @callback
+    def async_restore_last_state(self, last_state):
+        """Restore previous state."""
+        self._state = last_state.state == STATE_ON
+        if 'brightness' in last_state.attributes:
+            self._brightness = last_state.attributes['brightness']
+        if 'color_temp' in last_state.attributes:
+            self._color_temp = last_state.attributes['color_temp']
+        if 'hs_color' in last_state.attributes:
+            self._hs_color = last_state.attributes['hs_color']
+
     async def async_turn_on(self, **kwargs):
         """Turn the entity on."""
         transition = kwargs.get(light.ATTR_TRANSITION)
@@ -227,5 +240,10 @@ class Light(ZhaEntity, light.Light):
 
     async def async_update(self):
         """Attempt to retrieve on off state from the light."""
+        await super().async_update()
         if self._on_off_channel:
-            await self._on_off_channel.async_update()
+            self._state = await self._on_off_channel.get_attribute_value(
+                'on_off')
+        if self._level_channel:
+            self._brightness = await self._level_channel.get_attribute_value(
+                'current_level')

homeassistant/components/zha/sensor.py

@@ -6,8 +6,11 @@ at https://home-assistant.io/components/sensor.zha/
 """
 import logging
 
+from homeassistant.core import callback
 from homeassistant.components.sensor import DOMAIN
-from homeassistant.const import TEMP_CELSIUS, POWER_WATT
+from homeassistant.const import (
+    TEMP_CELSIUS, POWER_WATT, ATTR_UNIT_OF_MEASUREMENT
+)
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from .core.const import (
     DATA_ZHA, DATA_ZHA_DISPATCHERS, ZHA_DISCOVERY_NEW, HUMIDITY, TEMPERATURE,
@@ -133,22 +136,22 @@ class Sensor(ZhaEntity):
     def __init__(self, unique_id, zha_device, channels, **kwargs):
         """Init this sensor."""
         super().__init__(unique_id, zha_device, channels, **kwargs)
-        sensor_type = kwargs.get(SENSOR_TYPE, GENERIC)
+        self._sensor_type = kwargs.get(SENSOR_TYPE, GENERIC)
-        self._unit = UNIT_REGISTRY.get(sensor_type)
+        self._unit = UNIT_REGISTRY.get(self._sensor_type)
         self._formatter_function = FORMATTER_FUNC_REGISTRY.get(
-            sensor_type,
+            self._sensor_type,
             pass_through_formatter
         )
         self._force_update = FORCE_UPDATE_REGISTRY.get(
-            sensor_type,
+            self._sensor_type,
             False
         )
         self._should_poll = POLLING_REGISTRY.get(
-            sensor_type,
+            self._sensor_type,
             False
         )
         self._channel = self.cluster_channels.get(
-            CHANNEL_REGISTRY.get(sensor_type, ATTRIBUTE_CHANNEL)
+            CHANNEL_REGISTRY.get(self._sensor_type, ATTRIBUTE_CHANNEL)
         )
 
     async def async_added_to_hass(self):
@@ -176,5 +179,15 @@ class Sensor(ZhaEntity):
 
     def async_set_state(self, state):
         """Handle state update from channel."""
+        # this is necessary because HA saves the unit based on what shows in
+        # the UI and not based on what the sensor has configured so we need
+        # to flip it back after state restoration
+        self._unit = UNIT_REGISTRY.get(self._sensor_type)
         self._state = self._formatter_function(state)
         self.async_schedule_update_ha_state()
+
+    @callback
+    def async_restore_last_state(self, last_state):
+        """Restore previous state."""
+        self._state = last_state.state
+        self._unit = last_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)

homeassistant/components/zha/switch.py

@@ -7,6 +7,8 @@ at https://home-assistant.io/components/switch.zha/
 import logging
 
 from homeassistant.components.switch import DOMAIN, SwitchDevice
+from homeassistant.const import STATE_ON
+from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from .core.const import (
     DATA_ZHA, DATA_ZHA_DISPATCHERS, ZHA_DISCOVERY_NEW, ON_OFF_CHANNEL,
@@ -100,3 +102,15 @@ class Switch(ZhaEntity, SwitchDevice):
         await super().async_added_to_hass()
         await self.async_accept_signal(
             self._on_off_channel, SIGNAL_ATTR_UPDATED, self.async_set_state)
+
+    @callback
+    def async_restore_last_state(self, last_state):
+        """Restore previous state."""
+        self._state = last_state.state == STATE_ON
+
+    async def async_update(self):
+        """Attempt to retrieve on off state from the switch."""
+        await super().async_update()
+        if self._on_off_channel:
+            self._state = await self._on_off_channel.get_attribute_value(
+                'on_off')

tests/components/zha/conftest.py

@@ -10,6 +10,7 @@ from homeassistant.components.zha.core.gateway import establish_device_mappings
 from homeassistant.components.zha.core.channels.registry \
     import populate_channel_registry
 from .common import async_setup_entry
+from homeassistant.components.zha.core.store import async_get_registry
 
 
 @pytest.fixture(name='config_entry')
@@ -22,7 +23,7 @@ def config_entry_fixture(hass):
 
 
 @pytest.fixture(name='zha_gateway')
-def zha_gateway_fixture(hass):
+async def zha_gateway_fixture(hass):
     """Fixture representing a zha gateway.
 
     Create a ZHAGateway object that can be used to interact with as if we
@@ -34,7 +35,8 @@ def zha_gateway_fixture(hass):
         hass.data[DATA_ZHA][component] = (
             hass.data[DATA_ZHA].get(component, {})
         )
-    return ZHAGateway(hass, {})
+    zha_storage = await async_get_registry(hass)
+    return ZHAGateway(hass, {}, zha_storage)
 
 
 @pytest.fixture(autouse=True)