Fix migration of entities of Hue integration (#61095)

* fix device name in log

* Fix Hue migration for all id versions

* fix tests

* typo

* change to bit more universal approach

* fix test again

* formatting
pull/61115/head
Marcel van der Veldt 2021-12-06 17:24:59 +01:00 committed by GitHub
parent dddca8aaeb
commit bd8bba9e3f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 125 additions and 67 deletions

View File

@ -4,6 +4,7 @@ import logging
from aiohue import HueBridgeV2
from aiohue.discovery import is_v2_bridge
from aiohue.v2.models.device import DeviceArchetypes
from aiohue.v2.models.resource import ResourceTypes
from homeassistant import core
@ -18,7 +19,10 @@ from homeassistant.const import (
DEVICE_CLASS_TEMPERATURE,
)
from homeassistant.helpers import aiohttp_client
from homeassistant.helpers.device_registry import async_get as async_get_device_registry
from homeassistant.helpers.device_registry import (
async_entries_for_config_entry as devices_for_config_entries,
async_get as async_get_device_registry,
)
from homeassistant.helpers.entity_registry import (
async_entries_for_config_entry as entities_for_config_entry,
async_entries_for_device,
@ -82,6 +86,18 @@ async def handle_v2_migration(hass: core.HomeAssistant, entry: ConfigEntry) -> N
dev_reg = async_get_device_registry(hass)
ent_reg = async_get_entity_registry(hass)
LOGGER.info("Start of migration of devices and entities to support API schema 2")
# Create mapping of mac address to HA device id's.
# Identifier in dev reg should be mac-address,
# but in some cases it has a postfix like `-0b` or `-01`.
dev_ids = {}
for hass_dev in devices_for_config_entries(dev_reg, entry.entry_id):
for domain, mac in hass_dev.identifiers:
if domain != DOMAIN:
continue
normalized_mac = mac.split("-")[0]
dev_ids[normalized_mac] = hass_dev.id
# initialize bridge connection just for the migration
async with HueBridgeV2(host, api_key, websession) as api:
@ -92,83 +108,93 @@ async def handle_v2_migration(hass: core.HomeAssistant, entry: ConfigEntry) -> N
DEVICE_CLASS_TEMPERATURE: ResourceTypes.TEMPERATURE,
}
# handle entities attached to device
# migrate entities attached to a device
for hue_dev in api.devices:
zigbee = api.devices.get_zigbee_connectivity(hue_dev.id)
if not zigbee or not zigbee.mac_address:
# not a zigbee device or invalid mac
continue
# get/update existing device by V1 identifier (mac address)
# the device will now have both the old and the new identifier
identifiers = {(DOMAIN, hue_dev.id), (DOMAIN, zigbee.mac_address)}
hass_dev = dev_reg.async_get_or_create(
config_entry_id=entry.entry_id, identifiers=identifiers
)
LOGGER.info("Migrated device %s (%s)", hass_dev.name, hass_dev.id)
# loop through al entities for device and find match
for ent in async_entries_for_device(ent_reg, hass_dev.id, True):
# migrate light
if ent.entity_id.startswith("light"):
# should always return one lightid here
new_unique_id = next(iter(hue_dev.lights))
if ent.unique_id == new_unique_id:
continue # just in case
LOGGER.info(
"Migrating %s from unique id %s to %s",
ent.entity_id,
ent.unique_id,
new_unique_id,
)
ent_reg.async_update_entity(
ent.entity_id, new_unique_id=new_unique_id
)
continue
# migrate sensors
matched_dev_class = sensor_class_mapping.get(
ent.original_device_class or "unknown"
# get existing device by V1 identifier (mac address)
if hue_dev.product_data.product_archetype == DeviceArchetypes.BRIDGE_V2:
hass_dev_id = dev_ids.get(api.config.bridge_id.upper())
else:
hass_dev_id = dev_ids.get(zigbee.mac_address)
if hass_dev_id is None:
# can be safely ignored, this device does not exist in current config
LOGGER.debug(
"Ignoring device %s (%s) as it does not (yet) exist in the device registry",
hue_dev.metadata.name,
hue_dev.id,
)
if matched_dev_class is None:
continue
dev_reg.async_update_device(
hass_dev_id, new_identifiers={(DOMAIN, hue_dev.id)}
)
LOGGER.info("Migrated device %s (%s)", hue_dev.metadata.name, hass_dev_id)
# loop through all entities for device and find match
for ent in async_entries_for_device(ent_reg, hass_dev_id, True):
if ent.entity_id.startswith("light"):
# migrate light
# should always return one lightid here
new_unique_id = next(iter(hue_dev.lights), None)
else:
# migrate sensors
matched_dev_class = sensor_class_mapping.get(
ent.original_device_class or "unknown"
)
new_unique_id = next(
(
sensor.id
for sensor in api.devices.get_sensors(hue_dev.id)
if sensor.type == matched_dev_class
),
None,
)
if new_unique_id is None:
# this may happen if we're looking at orphaned or unsupported entity
LOGGER.warning(
"Skip migration of %s because it no longer exists on the bridge",
ent.entity_id,
)
continue
for sensor in api.devices.get_sensors(hue_dev.id):
if sensor.type != matched_dev_class:
continue
new_unique_id = sensor.id
if ent.unique_id == new_unique_id:
break # just in case
try:
ent_reg.async_update_entity(
ent.entity_id, new_unique_id=new_unique_id
)
except ValueError:
# assume edge case where the entity was already migrated in a previous run
# which got aborted somehow and we do not want
# to crash the entire integration init
LOGGER.warning(
"Skip migration of %s because it already exists",
ent.entity_id,
)
else:
LOGGER.info(
"Migrating %s from unique id %s to %s",
"Migrated entity %s from unique id %s to %s",
ent.entity_id,
ent.unique_id,
new_unique_id,
)
try:
ent_reg.async_update_entity(
ent.entity_id, new_unique_id=sensor.id
)
except ValueError:
# assume edge case where the entity was already migrated in a previous run
# which got aborted somehow and we do not want
# to crash the entire integration init
LOGGER.warning(
"Skip migration of %s because it already exists",
ent.entity_id,
)
break
# migrate entities that are not connected to a device (groups)
for ent in entities_for_config_entry(ent_reg, entry.entry_id):
if ent.device_id is not None:
continue
v1_id = f"/groups/{ent.unique_id}"
hue_group = api.groups.room.get_by_v1_id(v1_id)
if hue_group is None or hue_group.grouped_light is None:
# try again with zone
hue_group = api.groups.zone.get_by_v1_id(v1_id)
if "-" in ent.unique_id:
# handle case where unique id is v2-id of group/zone
hue_group = api.groups.get(ent.unique_id)
else:
# handle case where the unique id is just the v1 id
v1_id = f"/groups/{ent.unique_id}"
hue_group = api.groups.room.get_by_v1_id(
v1_id
) or api.groups.zone.get_by_v1_id(v1_id)
if hue_group is None or hue_group.grouped_light is None:
# this may happen if we're looking at some orphaned entity
LOGGER.warning(

View File

@ -54,12 +54,12 @@ async def test_light_entity_migration(
# create device/entity with V1 schema in registry
device = dev_reg.async_get_or_create(
config_entry_id=config_entry.entry_id,
identifiers={(hue.DOMAIN, "00:17:88:01:09:aa:bb:65")},
identifiers={(hue.DOMAIN, "00:17:88:01:09:aa:bb:65-0b")},
)
ent_reg.async_get_or_create(
"light",
hue.DOMAIN,
"00:17:88:01:09:aa:bb:65",
"00:17:88:01:09:aa:bb:65-0b",
suggested_object_id="migrated_light_1",
device_id=device.id,
)
@ -74,14 +74,13 @@ async def test_light_entity_migration(
):
await hue.migration.handle_v2_migration(hass, config_entry)
# migrated device should have new identifier (guid) and old style (mac)
# migrated device should now have the new identifier (guid) instead of old style (mac)
migrated_device = dev_reg.async_get(device.id)
assert migrated_device is not None
assert migrated_device.identifiers == {
(hue.DOMAIN, "0b216218-d811-4c95-8c55-bbcda50f9d50"),
(hue.DOMAIN, "00:17:88:01:09:aa:bb:65"),
(hue.DOMAIN, "0b216218-d811-4c95-8c55-bbcda50f9d50")
}
# the entity should have the new identifier (guid)
# the entity should have the new unique_id (guid)
migrated_entity = ent_reg.async_get("light.migrated_light_1")
assert migrated_entity is not None
assert migrated_entity.unique_id == "02cba059-9c2c-4d45-97e4-4f79b1bfbaa1"
@ -131,14 +130,13 @@ async def test_sensor_entity_migration(
):
await hue.migration.handle_v2_migration(hass, config_entry)
# migrated device should have new identifier (guid) and old style (mac)
# migrated device should now have the new identifier (guid) instead of old style (mac)
migrated_device = dev_reg.async_get(device.id)
assert migrated_device is not None
assert migrated_device.identifiers == {
(hue.DOMAIN, "2330b45d-6079-4c6e-bba6-1b68afb1a0d6"),
(hue.DOMAIN, device_mac),
(hue.DOMAIN, "2330b45d-6079-4c6e-bba6-1b68afb1a0d6")
}
# the entities should have the correct V2 identifier (guid)
# the entities should have the correct V2 unique_id (guid)
for dev_class, platform, new_id in sensor_mappings:
migrated_entity = ent_reg.async_get(
f"{platform}.hue_migrated_{dev_class}_sensor"
@ -147,7 +145,7 @@ async def test_sensor_entity_migration(
assert migrated_entity.unique_id == new_id
async def test_group_entity_migration(
async def test_group_entity_migration_with_v1_id(
hass, mock_bridge_v2, mock_config_entry_v2, v2_resources_test_data
):
"""Test if entity schema for grouped_lights migrates from v1 to v2."""
@ -156,6 +154,7 @@ async def test_group_entity_migration(
ent_reg = er.async_get(hass)
# create (deviceless) entity with V1 schema in registry
# using the legacy style group id as unique id
ent_reg.async_get_or_create(
"light",
hue.DOMAIN,
@ -177,3 +176,36 @@ async def test_group_entity_migration(
migrated_entity = ent_reg.async_get("light.hue_migrated_grouped_light")
assert migrated_entity is not None
assert migrated_entity.unique_id == "e937f8db-2f0e-49a0-936e-027e60e15b34"
async def test_group_entity_migration_with_v2_group_id(
hass, mock_bridge_v2, mock_config_entry_v2, v2_resources_test_data
):
"""Test if entity schema for grouped_lights migrates from v1 to v2."""
config_entry = mock_bridge_v2.config_entry = mock_config_entry_v2
ent_reg = er.async_get(hass)
# create (deviceless) entity with V1 schema in registry
# using the V2 group id as unique id
ent_reg.async_get_or_create(
"light",
hue.DOMAIN,
"6ddc9066-7e7d-4a03-a773-c73937968296",
suggested_object_id="hue_migrated_grouped_light",
config_entry=config_entry,
)
# now run the migration and check results
await mock_bridge_v2.api.load_test_data(v2_resources_test_data)
await hass.async_block_till_done()
with patch(
"homeassistant.components.hue.migration.HueBridgeV2",
return_value=mock_bridge_v2.api,
):
await hue.migration.handle_v2_migration(hass, config_entry)
# the entity should have the new identifier (guid)
migrated_entity = ent_reg.async_get("light.hue_migrated_grouped_light")
assert migrated_entity is not None
assert migrated_entity.unique_id == "e937f8db-2f0e-49a0-936e-027e60e15b34"