diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 03816c03df7..e7364e6665f 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -1,6 +1,7 @@ """Integration to UniFi controllers and its various features.""" from homeassistant.const import EVENT_HOMEASSISTANT_STOP from homeassistant.core import callback +from homeassistant.helpers import device_registry as dr from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC from .const import ( @@ -57,7 +58,7 @@ async def async_setup_entry(hass, config_entry): if controller.mac is None: return True - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, connections={(CONNECTION_NETWORK_MAC, controller.mac)}, diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index cea17e4e54c..b1ebbbe3475 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -39,7 +39,7 @@ from homeassistant.const import ( ) from homeassistant.core import callback from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady -from homeassistant.helpers import aiohttp_client +from homeassistant.helpers import aiohttp_client, entity_registry as er from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.entity_registry import async_entries_for_config_entry from homeassistant.helpers.event import async_track_time_interval @@ -333,7 +333,7 @@ class UniFiController: self._site_role = description[0]["site_role"] # Restore clients that are not a part of active clients list. - entity_registry = await self.hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(self.hass) for entry in async_entries_for_config_entry( entity_registry, self.config_entry.entry_id ): diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index faf51e6c853..b6e14717484 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -19,6 +19,7 @@ from homeassistant.components.device_tracker import DOMAIN from homeassistant.components.device_tracker.config_entry import ScannerEntity from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER from homeassistant.core import callback +from homeassistant.helpers import device_registry as dr from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC from homeassistant.helpers.dispatcher import async_dispatcher_connect import homeassistant.util.dt as dt_util @@ -417,8 +418,7 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity): async def async_update_device_registry(self) -> None: """Update device registry.""" - device_registry = await self.hass.helpers.device_registry.async_get_registry() - + device_registry = dr.async_get(self.hass) device_registry.async_get_or_create( config_entry_id=self.controller.config_entry.entry_id, **self.device_info ) diff --git a/homeassistant/components/unifi/services.py b/homeassistant/components/unifi/services.py index 10d297df883..17e894df913 100644 --- a/homeassistant/components/unifi/services.py +++ b/homeassistant/components/unifi/services.py @@ -4,6 +4,7 @@ import voluptuous as vol from homeassistant.const import ATTR_DEVICE_ID from homeassistant.core import callback +from homeassistant.helpers import device_registry as dr from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC from .const import DOMAIN as UNIFI_DOMAIN @@ -53,7 +54,7 @@ def async_unload_services(hass) -> None: async def async_reconnect_client(hass, data) -> None: """Try to get wireless client to reconnect to Wi-Fi.""" - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device_entry = device_registry.async_get(data[ATTR_DEVICE_ID]) mac = "" diff --git a/homeassistant/components/unifi/unifi_entity_base.py b/homeassistant/components/unifi/unifi_entity_base.py index 9d2d8071fca..e3b6e4f9970 100644 --- a/homeassistant/components/unifi/unifi_entity_base.py +++ b/homeassistant/components/unifi/unifi_entity_base.py @@ -3,6 +3,7 @@ import logging from typing import Any from homeassistant.core import callback +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import Entity from homeassistant.helpers.entity_registry import async_entries_for_device @@ -89,13 +90,13 @@ class UniFiBase(Entity): if self.key not in keys: return - entity_registry = await self.hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(self.hass) entity_entry = entity_registry.async_get(self.entity_id) if not entity_entry: await self.async_remove(force_remove=True) return - device_registry = await self.hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(self.hass) device_entry = device_registry.async_get(entity_entry.device_id) if not device_entry: entity_registry.async_remove(self.entity_id)