Clean up accessing device_registry helpers via hass (#72031)

This commit is contained in:
Franck Nijhof 2022-05-17 21:22:15 +02:00 committed by GitHub
parent 69e622b327
commit c3d19f3827
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
28 changed files with 105 additions and 87 deletions

View file

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

View file

@ -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

View file

@ -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(

View file

@ -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

View file

@ -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)

View file

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

View file

@ -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

View file

@ -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:

View file

@ -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:

View file

@ -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()

View file

@ -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"

View file

@ -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
)
)

View file

@ -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.

View file

@ -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)

View file

@ -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(

View file

@ -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(),

View file

@ -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",

View file

@ -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]

View file

@ -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
)

View file

@ -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]

View file

@ -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]

View file

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

View file

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

View file

@ -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)

View file

@ -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)

View file

@ -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."""

View file

@ -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"

View file

@ -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(