From cfaaae661acd54bd770a77a08a4aaa2263be9af2 Mon Sep 17 00:00:00 2001 From: Andrew Sayre <6730289+andrewsayre@users.noreply.github.com> Date: Tue, 30 Apr 2019 12:04:37 -0500 Subject: [PATCH] Add core APIs to migrate device identifiers and entity unique_id (#23481) * Add device identifiers migration * Add entity unique_id migration * Update per arch issue * Move to existing update methods --- homeassistant/helpers/device_registry.py | 10 +++++++-- homeassistant/helpers/entity_registry.py | 18 +++++++++++++--- tests/helpers/test_device_registry.py | 16 +++++++++++---- tests/helpers/test_entity_registry.py | 26 ++++++++++++++++++++++++ 4 files changed, 61 insertions(+), 9 deletions(-) diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py index 25c9933fd11..596bc84b6f9 100644 --- a/homeassistant/helpers/device_registry.py +++ b/homeassistant/helpers/device_registry.py @@ -134,16 +134,19 @@ class DeviceRegistry: @callback def async_update_device( - self, device_id, *, area_id=_UNDEF, name_by_user=_UNDEF): + self, device_id, *, area_id=_UNDEF, name_by_user=_UNDEF, + new_identifiers=_UNDEF): """Update properties of a device.""" return self._async_update_device( - device_id, area_id=area_id, name_by_user=name_by_user) + device_id, area_id=area_id, name_by_user=name_by_user, + new_identifiers=new_identifiers) @callback def _async_update_device(self, device_id, *, add_config_entry_id=_UNDEF, remove_config_entry_id=_UNDEF, merge_connections=_UNDEF, merge_identifiers=_UNDEF, + new_identifiers=_UNDEF, manufacturer=_UNDEF, model=_UNDEF, name=_UNDEF, @@ -178,6 +181,9 @@ class DeviceRegistry: if value is not _UNDEF and not value.issubset(old_value): changes[attr_name] = old_value | value + if new_identifiers is not _UNDEF: + changes['identifiers'] = new_identifiers + for attr_name, value in ( ('manufacturer', manufacturer), ('model', model), diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index be50d11d17d..64064ffde7b 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -160,18 +160,19 @@ class EntityRegistry: @callback def async_update_entity(self, entity_id, *, name=_UNDEF, - new_entity_id=_UNDEF): + new_entity_id=_UNDEF, new_unique_id=_UNDEF): """Update properties of an entity.""" return self._async_update_entity( entity_id, name=name, - new_entity_id=new_entity_id + new_entity_id=new_entity_id, + new_unique_id=new_unique_id ) @callback def _async_update_entity(self, entity_id, *, name=_UNDEF, config_entry_id=_UNDEF, new_entity_id=_UNDEF, - device_id=_UNDEF): + device_id=_UNDEF, new_unique_id=_UNDEF): """Private facing update properties method.""" old = self.entities[entity_id] @@ -201,6 +202,17 @@ class EntityRegistry: self.entities.pop(entity_id) entity_id = changes['entity_id'] = new_entity_id + if new_unique_id is not _UNDEF: + conflict = next((entity for entity in self.entities.values() + if entity.unique_id == new_unique_id + and entity.domain == old.domain + and entity.platform == old.platform), None) + if conflict: + raise ValueError( + "Unique id '{}' is already in use by '{}'".format( + new_unique_id, conflict.entity_id)) + changes['unique_id'] = new_unique_id + if not changes: return old diff --git a/tests/helpers/test_device_registry.py b/tests/helpers/test_device_registry.py index adfa05a021b..8c874a9837b 100644 --- a/tests/helpers/test_device_registry.py +++ b/tests/helpers/test_device_registry.py @@ -361,17 +361,25 @@ async def test_update(registry): config_entry_id='1234', connections={ (device_registry.CONNECTION_NETWORK_MAC, '12:34:56:AB:CD:EF') - }) - + }, + identifiers={('hue', '456'), ('bla', '123')}) + new_identifiers = { + ('hue', '654'), + ('bla', '321') + } assert not entry.area_id assert not entry.name_by_user - updated_entry = registry.async_update_device( - entry.id, area_id='12345A', name_by_user='Test Friendly Name') + with patch.object(registry, 'async_schedule_save') as mock_save: + updated_entry = registry.async_update_device( + entry.id, area_id='12345A', name_by_user='Test Friendly Name', + new_identifiers=new_identifiers) + assert mock_save.call_count == 1 assert updated_entry != entry assert updated_entry.area_id == '12345A' assert updated_entry.name_by_user == 'Test Friendly Name' + assert updated_entry.identifiers == new_identifiers async def test_loading_race_condition(hass): diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index 3fb79f693bd..529b03160ca 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -271,3 +271,29 @@ async def test_loading_race_condition(hass): mock_load.assert_called_once_with() assert results[0] == results[1] + + +async def test_update_entity_unique_id(registry): + """Test entity's unique_id is updated.""" + entry = registry.async_get_or_create( + 'light', 'hue', '5678', config_entry_id='mock-id-1') + new_unique_id = '1234' + with patch.object(registry, 'async_schedule_save') as mock_schedule_save: + updated_entry = registry.async_update_entity( + entry.entity_id, new_unique_id=new_unique_id) + assert updated_entry != entry + assert updated_entry.unique_id == new_unique_id + assert mock_schedule_save.call_count == 1 + + +async def test_update_entity_unique_id_conflict(registry): + """Test migration raises when unique_id already in use.""" + entry = registry.async_get_or_create( + 'light', 'hue', '5678', config_entry_id='mock-id-1') + entry2 = registry.async_get_or_create( + 'light', 'hue', '1234', config_entry_id='mock-id-1') + with patch.object(registry, 'async_schedule_save') as mock_schedule_save, \ + pytest.raises(ValueError): + registry.async_update_entity( + entry.entity_id, new_unique_id=entry2.unique_id) + assert mock_schedule_save.call_count == 0