Address late review of unifi

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)},

homeassistant/components/unifi/controller.py

@@ -39,7 +39,7 @@
 )
 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 @@ async def async_setup(self):
         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
         ):

homeassistant/components/unifi/device_tracker.py

@@ -19,6 +19,7 @@
 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 @@ def device_info(self):
 
     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
         )

homeassistant/components/unifi/services.py

@@ -4,6 +4,7 @@
 
 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 = ""

homeassistant/components/unifi/unifi_entity_base.py

@@ -3,6 +3,7 @@
 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 @@ async def remove_item(self, keys: set) -> None:
         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)