From c3d19f38274935ab8f784e6fbeb7c5959a551074 Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Tue, 17 May 2022 21:22:15 +0200 Subject: [PATCH] Clean up accessing device_registry helpers via hass (#72031) --- homeassistant/components/guardian/__init__.py | 2 +- homeassistant/components/hue/v1/hue_event.py | 5 ++--- homeassistant/components/insteon/api/aldb.py | 3 ++- .../components/insteon/api/device.py | 3 ++- homeassistant/components/insteon/utils.py | 3 ++- homeassistant/components/mikrotik/__init__.py | 4 ++-- homeassistant/components/nest/__init__.py | 8 +++++-- .../components/nest/device_trigger.py | 6 ++--- .../components/netatmo/device_trigger.py | 22 ++++++++++++++----- .../components/netatmo/netatmo_entity_base.py | 7 +++--- homeassistant/components/netatmo/sensor.py | 6 ++--- homeassistant/components/plex/__init__.py | 4 ++-- homeassistant/components/rfxtrx/__init__.py | 15 ++++--------- homeassistant/components/starline/__init__.py | 3 ++- .../components/tellduslive/__init__.py | 4 ++-- homeassistant/components/tradfri/__init__.py | 2 +- homeassistant/components/zha/__init__.py | 8 ++++--- homeassistant/components/zha/core/helpers.py | 3 ++- tests/components/deconz/test_deconz_event.py | 17 ++++++++------ tests/components/deconz/test_logbook.py | 5 +++-- .../home_plus_control/test_switch.py | 13 +++++------ tests/components/kaleidescape/test_init.py | 3 ++- .../kaleidescape/test_media_player.py | 3 ++- tests/components/kraken/test_sensor.py | 7 +++--- tests/components/mikrotik/test_init.py | 2 +- tests/components/picnic/test_sensor.py | 7 +++--- tests/components/unifi/test_init.py | 7 +++--- tests/components/unifi/test_services.py | 20 ++++++++--------- 28 files changed, 105 insertions(+), 87 deletions(-) diff --git a/homeassistant/components/guardian/__init__.py b/homeassistant/components/guardian/__init__.py index b971a428a76..25e0df913d8 100644 --- a/homeassistant/components/guardian/__init__.py +++ b/homeassistant/components/guardian/__init__.py @@ -363,7 +363,7 @@ class PairedSensorManager: # Remove the paired sensor device from the device registry (which will # clean up entities and the entity registry): - dev_reg = await self._hass.helpers.device_registry.async_get_registry() + dev_reg = dr.async_get(self._hass) device = dev_reg.async_get_or_create( config_entry_id=self._entry.entry_id, identifiers={(DOMAIN, uid)} ) diff --git a/homeassistant/components/hue/v1/hue_event.py b/homeassistant/components/hue/v1/hue_event.py index 9074baaaa88..b3faf88c2d9 100644 --- a/homeassistant/components/hue/v1/hue_event.py +++ b/homeassistant/components/hue/v1/hue_event.py @@ -10,6 +10,7 @@ from aiohue.v1.sensors import ( from homeassistant.const import CONF_DEVICE_ID, CONF_EVENT, CONF_ID, CONF_UNIQUE_ID from homeassistant.core import callback +from homeassistant.helpers import device_registry as dr from homeassistant.util import dt as dt_util, slugify from ..const import ATTR_HUE_EVENT @@ -89,9 +90,7 @@ class HueEvent(GenericHueDevice): async def async_update_device_registry(self): """Update device registry.""" - device_registry = ( - await self.bridge.hass.helpers.device_registry.async_get_registry() - ) + device_registry = dr.async_get(self.bridge.hass) entry = device_registry.async_get_or_create( config_entry_id=self.bridge.config_entry.entry_id, **self.device_info diff --git a/homeassistant/components/insteon/api/aldb.py b/homeassistant/components/insteon/api/aldb.py index a3132570ccc..a119707b03d 100644 --- a/homeassistant/components/insteon/api/aldb.py +++ b/homeassistant/components/insteon/api/aldb.py @@ -12,6 +12,7 @@ import voluptuous as vol from homeassistant.components import websocket_api from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers import device_registry as dr from ..const import DEVICE_ADDRESS, ID, INSTEON_DEVICE_NOT_FOUND, TYPE from .device import async_device_name, notify_device_not_found @@ -82,7 +83,7 @@ async def websocket_get_aldb( aldb.update(device.aldb.pending_changes) changed_records = list(device.aldb.pending_changes.keys()) - dev_registry = await hass.helpers.device_registry.async_get_registry() + dev_registry = dr.async_get(hass) records = [ await async_aldb_record_to_dict( diff --git a/homeassistant/components/insteon/api/device.py b/homeassistant/components/insteon/api/device.py index beef78394fa..ea6bade4835 100644 --- a/homeassistant/components/insteon/api/device.py +++ b/homeassistant/components/insteon/api/device.py @@ -6,6 +6,7 @@ import voluptuous as vol from homeassistant.components import websocket_api from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers import device_registry as dr from ..const import ( DEVICE_ADDRESS, @@ -68,7 +69,7 @@ async def websocket_get_device( msg: dict, ) -> None: """Get an Insteon device.""" - dev_registry = await hass.helpers.device_registry.async_get_registry() + dev_registry = dr.async_get(hass) if not (ha_device := dev_registry.async_get(msg[DEVICE_ID])): notify_device_not_found(connection, msg, HA_DEVICE_NOT_FOUND) return diff --git a/homeassistant/components/insteon/utils.py b/homeassistant/components/insteon/utils.py index 03647559345..e8e34ee8e62 100644 --- a/homeassistant/components/insteon/utils.py +++ b/homeassistant/components/insteon/utils.py @@ -28,6 +28,7 @@ from homeassistant.const import ( ENTITY_MATCH_ALL, ) from homeassistant.core import ServiceCall, callback +from homeassistant.helpers import device_registry as dr from homeassistant.helpers.dispatcher import ( async_dispatcher_connect, async_dispatcher_send, @@ -294,7 +295,7 @@ def async_register_services(hass): """Remove the device and all entities from hass.""" signal = f"{address.id}_{SIGNAL_REMOVE_ENTITY}" async_dispatcher_send(hass, signal) - dev_registry = await hass.helpers.device_registry.async_get_registry() + dev_registry = dr.async_get(hass) device = dev_registry.async_get_device(identifiers={(DOMAIN, str(address))}) if device: dev_registry.async_remove_device(device.id) diff --git a/homeassistant/components/mikrotik/__init__.py b/homeassistant/components/mikrotik/__init__.py index f0dc249d26a..1ef250a3f4e 100644 --- a/homeassistant/components/mikrotik/__init__.py +++ b/homeassistant/components/mikrotik/__init__.py @@ -11,7 +11,7 @@ from homeassistant.const import ( CONF_VERIFY_SSL, ) from homeassistant.core import HomeAssistant -from homeassistant.helpers import config_validation as cv +from homeassistant.helpers import config_validation as cv, device_registry as dr from homeassistant.helpers.typing import ConfigType from .const import ( @@ -76,7 +76,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b return False hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = hub - 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={(DOMAIN, hub.serial_num)}, diff --git a/homeassistant/components/nest/__init__.py b/homeassistant/components/nest/__init__.py index 5a34df1d74b..7d8353738ff 100644 --- a/homeassistant/components/nest/__init__.py +++ b/homeassistant/components/nest/__init__.py @@ -42,7 +42,11 @@ from homeassistant.exceptions import ( HomeAssistantError, Unauthorized, ) -from homeassistant.helpers import config_validation as cv, entity_registry as er +from homeassistant.helpers import ( + config_validation as cv, + device_registry as dr, + entity_registry as er, +) from homeassistant.helpers.entity_registry import async_entries_for_device from homeassistant.helpers.typing import ConfigType @@ -145,7 +149,7 @@ class SignalUpdateCallback: if not (events := event_message.resource_update_events): return _LOGGER.debug("Event Update %s", events.keys()) - device_registry = await self._hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(self._hass) device_entry = device_registry.async_get_device({(DOMAIN, device_id)}) if not device_entry: return diff --git a/homeassistant/components/nest/device_trigger.py b/homeassistant/components/nest/device_trigger.py index 28eb444b91a..67b361e76d8 100644 --- a/homeassistant/components/nest/device_trigger.py +++ b/homeassistant/components/nest/device_trigger.py @@ -16,7 +16,7 @@ from homeassistant.components.device_automation.exceptions import ( from homeassistant.components.homeassistant.triggers import event as event_trigger from homeassistant.const import CONF_DEVICE_ID, CONF_DOMAIN, CONF_PLATFORM, CONF_TYPE from homeassistant.core import CALLBACK_TYPE, HomeAssistant -from homeassistant.helpers.device_registry import DeviceRegistry +from homeassistant.helpers import device_registry as dr from homeassistant.helpers.typing import ConfigType from .const import DATA_SUBSCRIBER, DOMAIN @@ -35,9 +35,7 @@ TRIGGER_SCHEMA = DEVICE_TRIGGER_BASE_SCHEMA.extend( async def async_get_nest_device_id(hass: HomeAssistant, device_id: str) -> str | None: """Get the nest API device_id from the HomeAssistant device_id.""" - device_registry: DeviceRegistry = ( - await hass.helpers.device_registry.async_get_registry() - ) + device_registry = dr.async_get(hass) if device := device_registry.async_get(device_id): for (domain, unique_id) in device.identifiers: if domain == DOMAIN: diff --git a/homeassistant/components/netatmo/device_trigger.py b/homeassistant/components/netatmo/device_trigger.py index 7a9acac2f96..ae62956b691 100644 --- a/homeassistant/components/netatmo/device_trigger.py +++ b/homeassistant/components/netatmo/device_trigger.py @@ -23,7 +23,11 @@ from homeassistant.const import ( CONF_TYPE, ) from homeassistant.core import CALLBACK_TYPE, HomeAssistant -from homeassistant.helpers import config_validation as cv, entity_registry +from homeassistant.helpers import ( + config_validation as cv, + device_registry as dr, + entity_registry, +) from homeassistant.helpers.typing import ConfigType from .climate import STATE_NETATMO_AWAY, STATE_NETATMO_HG, STATE_NETATMO_SCHEDULE @@ -74,9 +78,14 @@ async def async_validate_trigger_config( """Validate config.""" config = TRIGGER_SCHEMA(config) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device = device_registry.async_get(config[CONF_DEVICE_ID]) + if not device: + raise InvalidDeviceAutomationConfig( + f"Trigger invalid, device with ID {config[CONF_DEVICE_ID]} not found" + ) + trigger = config[CONF_TYPE] if ( @@ -94,11 +103,14 @@ async def async_get_triggers( ) -> list[dict[str, Any]]: """List device triggers for Netatmo devices.""" registry = entity_registry.async_get(hass) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) triggers = [] for entry in entity_registry.async_entries_for_device(registry, device_id): - device = device_registry.async_get(device_id) + if ( + device := device_registry.async_get(device_id) + ) is None or device.model is None: + continue for trigger in DEVICES.get(device.model, []): if trigger in SUBTYPES: @@ -134,7 +146,7 @@ async def async_attach_trigger( automation_info: AutomationTriggerInfo, ) -> CALLBACK_TYPE: """Attach a trigger.""" - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device = device_registry.async_get(config[CONF_DEVICE_ID]) if not device: diff --git a/homeassistant/components/netatmo/netatmo_entity_base.py b/homeassistant/components/netatmo/netatmo_entity_base.py index 56f25e04906..decedbbdfbd 100644 --- a/homeassistant/components/netatmo/netatmo_entity_base.py +++ b/homeassistant/components/netatmo/netatmo_entity_base.py @@ -3,6 +3,7 @@ from __future__ import annotations from homeassistant.const import ATTR_ATTRIBUTION from homeassistant.core import callback +from homeassistant.helpers import device_registry as dr from homeassistant.helpers.entity import DeviceInfo, Entity from .const import ( @@ -65,9 +66,9 @@ class NetatmoBase(Entity): if sub is None: await self.data_handler.unregister_data_class(signal_name, None) - registry = await self.hass.helpers.device_registry.async_get_registry() - device = registry.async_get_device({(DOMAIN, self._id)}) - self.hass.data[DOMAIN][DATA_DEVICE_IDS][self._id] = device.id + registry = dr.async_get(self.hass) + if device := registry.async_get_device({(DOMAIN, self._id)}): + self.hass.data[DOMAIN][DATA_DEVICE_IDS][self._id] = device.id self.async_update_callback() diff --git a/homeassistant/components/netatmo/sensor.py b/homeassistant/components/netatmo/sensor.py index 41ae27b2992..217b2146cc9 100644 --- a/homeassistant/components/netatmo/sensor.py +++ b/homeassistant/components/netatmo/sensor.py @@ -29,7 +29,7 @@ from homeassistant.const import ( ) from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import PlatformNotReady -from homeassistant.helpers.device_registry import async_entries_for_config_entry +from homeassistant.helpers import device_registry as dr from homeassistant.helpers.dispatcher import ( async_dispatcher_connect, async_dispatcher_send, @@ -393,13 +393,13 @@ async def async_setup_entry( async_add_entities(await find_entities(data_class_name), True) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) async def add_public_entities(update: bool = True) -> None: """Retrieve Netatmo public weather entities.""" entities = { device.name: device.id - for device in async_entries_for_config_entry( + for device in dr.async_entries_for_config_entry( device_registry, entry.entry_id ) if device.model == "Public Weather stations" diff --git a/homeassistant/components/plex/__init__.py b/homeassistant/components/plex/__init__.py index 4ca83d98242..dbfe55077d7 100644 --- a/homeassistant/components/plex/__init__.py +++ b/homeassistant/components/plex/__init__.py @@ -299,14 +299,14 @@ def async_cleanup_plex_devices(hass, entry): device_registry = dev_reg.async_get(hass) entity_registry = ent_reg.async_get(hass) - device_entries = hass.helpers.device_registry.async_entries_for_config_entry( + device_entries = dev_reg.async_entries_for_config_entry( device_registry, entry.entry_id ) for device_entry in device_entries: if ( len( - hass.helpers.entity_registry.async_entries_for_device( + ent_reg.async_entries_for_device( entity_registry, device_entry.id, include_disabled_entities=True ) ) diff --git a/homeassistant/components/rfxtrx/__init__.py b/homeassistant/components/rfxtrx/__init__.py index 89521aba9f2..b517abafc86 100644 --- a/homeassistant/components/rfxtrx/__init__.py +++ b/homeassistant/components/rfxtrx/__init__.py @@ -24,12 +24,7 @@ from homeassistant.const import ( Platform, ) from homeassistant.core import Event, HomeAssistant, ServiceCall, callback -import homeassistant.helpers.config_validation as cv -from homeassistant.helpers.device_registry import ( - EVENT_DEVICE_REGISTRY_UPDATED, - DeviceEntry, - DeviceRegistry, -) +from homeassistant.helpers import config_validation as cv, device_registry as dr from homeassistant.helpers.dispatcher import ( async_dispatcher_connect, async_dispatcher_send, @@ -174,9 +169,7 @@ async def async_setup_internal(hass, entry: ConfigEntry): devices = _get_device_lookup(config[CONF_DEVICES]) pt2262_devices: list[str] = [] - device_registry: DeviceRegistry = ( - await hass.helpers.device_registry.async_get_registry() - ) + device_registry = dr.async_get(hass) # Declare the Handle event @callback @@ -268,7 +261,7 @@ async def async_setup_internal(hass, entry: ConfigEntry): _remove_device(device_id) entry.async_on_unload( - hass.bus.async_listen(EVENT_DEVICE_REGISTRY_UPDATED, _updated_device) + hass.bus.async_listen(dr.EVENT_DEVICE_REGISTRY_UPDATED, _updated_device) ) def _shutdown_rfxtrx(event): @@ -448,7 +441,7 @@ def get_device_tuple_from_identifiers( async def async_remove_config_entry_device( - hass: HomeAssistant, config_entry: ConfigEntry, device_entry: DeviceEntry + hass: HomeAssistant, config_entry: ConfigEntry, device_entry: dr.DeviceEntry ) -> bool: """Remove config entry from a device. diff --git a/homeassistant/components/starline/__init__.py b/homeassistant/components/starline/__init__.py index 7d05c0d47ee..886a583ae93 100644 --- a/homeassistant/components/starline/__init__.py +++ b/homeassistant/components/starline/__init__.py @@ -7,6 +7,7 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_SCAN_INTERVAL from homeassistant.core import HomeAssistant, ServiceCall from homeassistant.exceptions import ConfigEntryNotReady +from homeassistant.helpers import device_registry as dr from .account import StarlineAccount from .const import ( @@ -33,7 +34,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hass.data[DOMAIN] = {} hass.data[DOMAIN][entry.entry_id] = account - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) for device in account.api.devices.values(): device_registry.async_get_or_create( config_entry_id=entry.entry_id, **account.device_info(device) diff --git a/homeassistant/components/tellduslive/__init__.py b/homeassistant/components/tellduslive/__init__.py index 572950b816b..92feb69c60c 100644 --- a/homeassistant/components/tellduslive/__init__.py +++ b/homeassistant/components/tellduslive/__init__.py @@ -10,7 +10,7 @@ from homeassistant import config_entries from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_HOST, CONF_SCAN_INTERVAL from homeassistant.core import HomeAssistant -import homeassistant.helpers.config_validation as cv +from homeassistant.helpers import config_validation as cv, device_registry as dr from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.event import async_call_later from homeassistant.helpers.typing import ConfigType @@ -84,7 +84,7 @@ async def async_new_client(hass, session, entry): _LOGGER.debug("Update interval %s seconds", interval) client = TelldusLiveClient(hass, entry, session, interval) hass.data[DOMAIN] = client - dev_reg = await hass.helpers.device_registry.async_get_registry() + dev_reg = dr.async_get(hass) for hub in await client.async_get_hubs(): _LOGGER.debug("Connected hub %s", hub["name"]) dev_reg.async_get_or_create( diff --git a/homeassistant/components/tradfri/__init__.py b/homeassistant/components/tradfri/__init__.py index e4d13568e6d..e79cd8396e1 100644 --- a/homeassistant/components/tradfri/__init__.py +++ b/homeassistant/components/tradfri/__init__.py @@ -84,7 +84,7 @@ async def async_setup_entry( await factory.shutdown() raise ConfigEntryNotReady from exc - dev_reg = await hass.helpers.device_registry.async_get_registry() + dev_reg = dr.async_get(hass) dev_reg.async_get_or_create( config_entry_id=entry.entry_id, connections=set(), diff --git a/homeassistant/components/zha/__init__.py b/homeassistant/components/zha/__init__.py index c512e104ae8..0e11d992a25 100644 --- a/homeassistant/components/zha/__init__.py +++ b/homeassistant/components/zha/__init__.py @@ -9,8 +9,8 @@ from zigpy.config import CONF_DEVICE, CONF_DEVICE_PATH from homeassistant import const as ha_const from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant +from homeassistant.helpers import device_registry as dr import homeassistant.helpers.config_validation as cv -from homeassistant.helpers.device_registry import CONNECTION_ZIGBEE from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.typing import ConfigType @@ -106,10 +106,12 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b coro = hass.config_entries.async_forward_entry_setup(config_entry, platform) zha_data[DATA_ZHA_PLATFORM_LOADED].append(hass.async_create_task(coro)) - 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_ZIGBEE, str(zha_gateway.application_controller.ieee))}, + connections={ + (dr.CONNECTION_ZIGBEE, str(zha_gateway.application_controller.ieee)) + }, identifiers={(DOMAIN, str(zha_gateway.application_controller.ieee))}, name="Zigbee Coordinator", manufacturer="ZHA", diff --git a/homeassistant/components/zha/core/helpers.py b/homeassistant/components/zha/core/helpers.py index fcd29c1619f..19a1d3fef9a 100644 --- a/homeassistant/components/zha/core/helpers.py +++ b/homeassistant/components/zha/core/helpers.py @@ -25,6 +25,7 @@ import zigpy.zdo.types as zdo_types from homeassistant.config_entries import ConfigEntry from homeassistant.core import State, callback +from homeassistant.helpers import device_registry as dr from .const import ( CLUSTER_TYPE_IN, @@ -161,7 +162,7 @@ def async_cluster_exists(hass, cluster_id): async def async_get_zha_device(hass, device_id): """Get a ZHA device for the given device registry id.""" - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) registry_device = device_registry.async_get(device_id) zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] ieee_address = list(list(registry_device.identifiers)[0])[1] diff --git a/tests/components/deconz/test_deconz_event.py b/tests/components/deconz/test_deconz_event.py index 5b04b459278..e697edd5a9a 100644 --- a/tests/components/deconz/test_deconz_event.py +++ b/tests/components/deconz/test_deconz_event.py @@ -22,7 +22,7 @@ from homeassistant.const import ( CONF_UNIQUE_ID, STATE_UNAVAILABLE, ) -from homeassistant.helpers.device_registry import async_entries_for_config_entry +from homeassistant.helpers import device_registry as dr from .test_gateway import DECONZ_WEB_REQUEST, setup_deconz_integration @@ -73,12 +73,13 @@ async def test_deconz_events(hass, aioclient_mock, mock_deconz_websocket): with patch.dict(DECONZ_WEB_REQUEST, data): config_entry = await setup_deconz_integration(hass, aioclient_mock) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) assert len(hass.states.async_all()) == 3 # 5 switches + 2 additional devices for deconz service and host assert ( - len(async_entries_for_config_entry(device_registry, config_entry.entry_id)) == 7 + len(dr.async_entries_for_config_entry(device_registry, config_entry.entry_id)) + == 7 ) assert hass.states.get("sensor.switch_2_battery").state == "100" assert hass.states.get("sensor.switch_3_battery").state == "100" @@ -267,12 +268,13 @@ async def test_deconz_alarm_events(hass, aioclient_mock, mock_deconz_websocket): with patch.dict(DECONZ_WEB_REQUEST, data): config_entry = await setup_deconz_integration(hass, aioclient_mock) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) assert len(hass.states.async_all()) == 4 # 1 alarm control device + 2 additional devices for deconz service and host assert ( - len(async_entries_for_config_entry(device_registry, config_entry.entry_id)) == 3 + len(dr.async_entries_for_config_entry(device_registry, config_entry.entry_id)) + == 3 ) captured_events = async_capture_events(hass, CONF_DECONZ_ALARM_EVENT) @@ -435,9 +437,10 @@ async def test_deconz_events_bad_unique_id(hass, aioclient_mock, mock_deconz_web with patch.dict(DECONZ_WEB_REQUEST, data): config_entry = await setup_deconz_integration(hass, aioclient_mock) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) assert len(hass.states.async_all()) == 1 assert ( - len(async_entries_for_config_entry(device_registry, config_entry.entry_id)) == 2 + len(dr.async_entries_for_config_entry(device_registry, config_entry.entry_id)) + == 2 ) diff --git a/tests/components/deconz/test_logbook.py b/tests/components/deconz/test_logbook.py index a78e795573a..9ba0799d04e 100644 --- a/tests/components/deconz/test_logbook.py +++ b/tests/components/deconz/test_logbook.py @@ -15,6 +15,7 @@ from homeassistant.const import ( CONF_UNIQUE_ID, STATE_ALARM_ARMED_AWAY, ) +from homeassistant.helpers import device_registry as dr from homeassistant.setup import async_setup_component from homeassistant.util import slugify @@ -56,7 +57,7 @@ async def test_humanifying_deconz_alarm_event(hass, aioclient_mock): with patch.dict(DECONZ_WEB_REQUEST, data): await setup_deconz_integration(hass, aioclient_mock) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) keypad_event_id = slugify(data["sensors"]["1"]["name"]) keypad_serial = data["sensors"]["1"]["uniqueid"].split("-", 1)[0] @@ -127,7 +128,7 @@ async def test_humanifying_deconz_event(hass, aioclient_mock): with patch.dict(DECONZ_WEB_REQUEST, data): await setup_deconz_integration(hass, aioclient_mock) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) switch_event_id = slugify(data["sensors"]["1"]["name"]) switch_serial = data["sensors"]["1"]["uniqueid"].split("-", 1)[0] diff --git a/tests/components/home_plus_control/test_switch.py b/tests/components/home_plus_control/test_switch.py index 75d416ba2b1..44a969392bf 100644 --- a/tests/components/home_plus_control/test_switch.py +++ b/tests/components/home_plus_control/test_switch.py @@ -16,6 +16,7 @@ from homeassistant.const import ( STATE_ON, STATE_UNAVAILABLE, ) +from homeassistant.helpers import device_registry as dr, entity_registry as er from tests.common import async_fire_time_changed from tests.components.home_plus_control.conftest import ( @@ -33,8 +34,8 @@ def entity_assertions( expected_devices=None, ): """Assert number of entities and devices.""" - entity_reg = hass.helpers.entity_registry.async_get(hass) - device_reg = hass.helpers.device_registry.async_get(hass) + entity_reg = er.async_get(hass) + device_reg = dr.async_get(hass) if num_exp_devices is None: num_exp_devices = num_exp_entities @@ -53,13 +54,11 @@ def entity_assertions( def one_entity_state(hass, device_uid): """Assert the presence of an entity and return its state.""" - entity_reg = hass.helpers.entity_registry.async_get(hass) - device_reg = hass.helpers.device_registry.async_get(hass) + entity_reg = er.async_get(hass) + device_reg = dr.async_get(hass) device_id = device_reg.async_get_device({(DOMAIN, device_uid)}).id - entity_entries = hass.helpers.entity_registry.async_entries_for_device( - entity_reg, device_id - ) + entity_entries = er.async_entries_for_device(entity_reg, device_id) assert len(entity_entries) == 1 entity_entry = entity_entries[0] diff --git a/tests/components/kaleidescape/test_init.py b/tests/components/kaleidescape/test_init.py index 876c02ba5a6..d0826f4714a 100644 --- a/tests/components/kaleidescape/test_init.py +++ b/tests/components/kaleidescape/test_init.py @@ -5,6 +5,7 @@ from unittest.mock import AsyncMock from homeassistant.components.kaleidescape.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant +from homeassistant.helpers import device_registry as dr from . import MOCK_SERIAL @@ -50,7 +51,7 @@ async def test_device( mock_integration: MockConfigEntry, ) -> None: """Test device.""" - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device = device_registry.async_get_device( identifiers={("kaleidescape", MOCK_SERIAL)} ) diff --git a/tests/components/kaleidescape/test_media_player.py b/tests/components/kaleidescape/test_media_player.py index 94ba7f82fe8..11f5d5f5f2b 100644 --- a/tests/components/kaleidescape/test_media_player.py +++ b/tests/components/kaleidescape/test_media_player.py @@ -21,6 +21,7 @@ from homeassistant.const import ( STATE_PLAYING, ) from homeassistant.core import HomeAssistant +from homeassistant.helpers import device_registry as dr from . import MOCK_SERIAL @@ -173,7 +174,7 @@ async def test_device( mock_integration: MockConfigEntry, ) -> None: """Test device attributes.""" - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device = device_registry.async_get_device( identifiers={("kaleidescape", MOCK_SERIAL)} ) diff --git a/tests/components/kraken/test_sensor.py b/tests/components/kraken/test_sensor.py index cbfda938f99..1d6b9e6518d 100644 --- a/tests/components/kraken/test_sensor.py +++ b/tests/components/kraken/test_sensor.py @@ -11,8 +11,7 @@ from homeassistant.components.kraken.const import ( DOMAIN, ) from homeassistant.const import CONF_SCAN_INTERVAL, EVENT_HOMEASSISTANT_START -from homeassistant.helpers import entity_registry as er -from homeassistant.helpers.device_registry import DeviceEntryType +from homeassistant.helpers import device_registry as dr, entity_registry as er import homeassistant.util.dt as dt_util from .const import ( @@ -249,13 +248,13 @@ async def test_sensors_available_after_restart(hass): }, ) - 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=entry.entry_id, identifiers={(DOMAIN, "XBT_USD")}, name="XBT USD", manufacturer="Kraken.com", - entry_type=DeviceEntryType.SERVICE, + entry_type=dr.DeviceEntryType.SERVICE, ) entry.add_to_hass(hass) diff --git a/tests/components/mikrotik/test_init.py b/tests/components/mikrotik/test_init.py index 281b70e36be..30fa1a0a89f 100644 --- a/tests/components/mikrotik/test_init.py +++ b/tests/components/mikrotik/test_init.py @@ -25,7 +25,7 @@ async def test_successful_config_entry(hass): mock_registry = Mock() with patch.object(mikrotik, "MikrotikHub") as mock_hub, patch( - "homeassistant.helpers.device_registry.async_get_registry", + "homeassistant.components.mikrotik.dr.async_get", return_value=mock_registry, ): mock_hub.return_value.async_setup = AsyncMock(return_value=True) diff --git a/tests/components/picnic/test_sensor.py b/tests/components/picnic/test_sensor.py index 808f1ce6f41..8d680327bc8 100644 --- a/tests/components/picnic/test_sensor.py +++ b/tests/components/picnic/test_sensor.py @@ -17,8 +17,7 @@ from homeassistant.const import ( STATE_UNAVAILABLE, STATE_UNKNOWN, ) -from homeassistant.helpers import entity_registry as er -from homeassistant.helpers.device_registry import DeviceEntryType +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.util import dt from tests.common import ( @@ -497,13 +496,13 @@ class TestPicnicSensor(unittest.IsolatedAsyncioTestCase): # Setup platform and default mock responses await self._setup_platform(use_default_responses=True) - device_registry = await self.hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(self.hass) picnic_service = device_registry.async_get_device( identifiers={(const.DOMAIN, DEFAULT_USER_RESPONSE["user_id"])} ) assert picnic_service.model == DEFAULT_USER_RESPONSE["user_id"] assert picnic_service.name == "Picnic: Commonstreet 123a" - assert picnic_service.entry_type is DeviceEntryType.SERVICE + assert picnic_service.entry_type is dr.DeviceEntryType.SERVICE async def test_auth_token_is_saved_on_update(self): """Test that auth-token changes in the session object are reflected by the config entry.""" diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index b2d37ad7ee3..f183e1c22ff 100644 --- a/tests/components/unifi/test_init.py +++ b/tests/components/unifi/test_init.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, patch from homeassistant.components import unifi from homeassistant.components.unifi import async_flatten_entry_data from homeassistant.components.unifi.const import CONF_CONTROLLER, DOMAIN as UNIFI_DOMAIN -from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC +from homeassistant.helpers import device_registry as dr from homeassistant.setup import async_setup_component from .test_controller import ( @@ -53,9 +53,10 @@ async def test_controller_mac(hass): assert len(mock_controller.mock_calls) == 2 - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device = device_registry.async_get_or_create( - config_entry_id=entry.entry_id, connections={(CONNECTION_NETWORK_MAC, "mac1")} + config_entry_id=entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "mac1")}, ) assert device.configuration_url == "https://123:443" assert device.manufacturer == "Ubiquiti Networks" diff --git a/tests/components/unifi/test_services.py b/tests/components/unifi/test_services.py index 27e4ddea930..0c6f20869c8 100644 --- a/tests/components/unifi/test_services.py +++ b/tests/components/unifi/test_services.py @@ -9,7 +9,7 @@ from homeassistant.components.unifi.services import ( SUPPORTED_SERVICES, ) from homeassistant.const import ATTR_DEVICE_ID -from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC +from homeassistant.helpers import device_registry as dr from .test_controller import setup_unifi_integration @@ -62,10 +62,10 @@ async def test_reconnect_client(hass, aioclient_mock): f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr", ) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, - connections={(CONNECTION_NETWORK_MAC, clients[0]["mac"])}, + connections={(dr.CONNECTION_NETWORK_MAC, clients[0]["mac"])}, ) await hass.services.async_call( @@ -98,7 +98,7 @@ async def test_reconnect_device_without_mac(hass, aioclient_mock): aioclient_mock.clear_requests() - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, connections={("other connection", "not mac")}, @@ -132,10 +132,10 @@ async def test_reconnect_client_controller_unavailable(hass, aioclient_mock): f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr", ) - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, - connections={(CONNECTION_NETWORK_MAC, clients[0]["mac"])}, + connections={(dr.CONNECTION_NETWORK_MAC, clients[0]["mac"])}, ) await hass.services.async_call( @@ -153,10 +153,10 @@ async def test_reconnect_client_unknown_mac(hass, aioclient_mock): aioclient_mock.clear_requests() - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, - connections={(CONNECTION_NETWORK_MAC, "mac unknown to controller")}, + connections={(dr.CONNECTION_NETWORK_MAC, "mac unknown to controller")}, ) await hass.services.async_call( @@ -182,10 +182,10 @@ async def test_reconnect_wired_client(hass, aioclient_mock): aioclient_mock.clear_requests() - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, - connections={(CONNECTION_NETWORK_MAC, clients[0]["mac"])}, + connections={(dr.CONNECTION_NETWORK_MAC, clients[0]["mac"])}, ) await hass.services.async_call(