From 691d49f23ba42c4ee7f053920b10b41322913730 Mon Sep 17 00:00:00 2001 From: Christopher Bailey Date: Sat, 18 Jun 2022 13:56:28 -0400 Subject: [PATCH] Refactor migration code for UniFi Protect (#73499) --- .../components/unifiprotect/__init__.py | 60 +------ .../components/unifiprotect/migrate.py | 83 +++++++++ tests/components/unifiprotect/test_init.py | 144 ---------------- tests/components/unifiprotect/test_migrate.py | 158 ++++++++++++++++++ 4 files changed, 244 insertions(+), 201 deletions(-) create mode 100644 homeassistant/components/unifiprotect/migrate.py create mode 100644 tests/components/unifiprotect/test_migrate.py diff --git a/homeassistant/components/unifiprotect/__init__.py b/homeassistant/components/unifiprotect/__init__.py index 4ec11a899e3..d05f544ada1 100644 --- a/homeassistant/components/unifiprotect/__init__.py +++ b/homeassistant/components/unifiprotect/__init__.py @@ -17,11 +17,10 @@ from homeassistant.const import ( CONF_USERNAME, CONF_VERIFY_SSL, EVENT_HOMEASSISTANT_STOP, - Platform, ) from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady -from homeassistant.helpers import device_registry as dr, entity_registry as er +from homeassistant.helpers import device_registry as dr from homeassistant.helpers.aiohttp_client import async_create_clientsession from .const import ( @@ -37,6 +36,7 @@ from .const import ( ) from .data import ProtectData, async_ufp_instance_for_config_entry_ids from .discovery import async_start_discovery +from .migrate import async_migrate_data from .services import async_cleanup_services, async_setup_services from .utils import _async_unifi_mac_from_hass, async_get_devices @@ -45,60 +45,6 @@ _LOGGER = logging.getLogger(__name__) SCAN_INTERVAL = timedelta(seconds=DEFAULT_SCAN_INTERVAL) -async def _async_migrate_data( - hass: HomeAssistant, entry: ConfigEntry, protect: ProtectApiClient -) -> None: - - registry = er.async_get(hass) - to_migrate = [] - for entity in er.async_entries_for_config_entry(registry, entry.entry_id): - if entity.domain == Platform.BUTTON and "_" not in entity.unique_id: - _LOGGER.debug("Button %s needs migration", entity.entity_id) - to_migrate.append(entity) - - if len(to_migrate) == 0: - _LOGGER.debug("No entities need migration") - return - - _LOGGER.info("Migrating %s reboot button entities ", len(to_migrate)) - bootstrap = await protect.get_bootstrap() - count = 0 - for button in to_migrate: - device = None - for model in DEVICES_THAT_ADOPT: - attr = f"{model.value}s" - device = getattr(bootstrap, attr).get(button.unique_id) - if device is not None: - break - - if device is None: - continue - - new_unique_id = f"{device.id}_reboot" - _LOGGER.debug( - "Migrating entity %s (old unique_id: %s, new unique_id: %s)", - button.entity_id, - button.unique_id, - new_unique_id, - ) - try: - registry.async_update_entity(button.entity_id, new_unique_id=new_unique_id) - except ValueError: - _LOGGER.warning( - "Could not migrate entity %s (old unique_id: %s, new unique_id: %s)", - button.entity_id, - button.unique_id, - new_unique_id, - ) - else: - count += 1 - - if count < len(to_migrate): - _LOGGER.warning("Failed to migate %s reboot buttons", len(to_migrate) - count) - else: - _LOGGER.info("Migrated %s reboot button entities", count) - - async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up the UniFi Protect config entries.""" @@ -133,7 +79,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ) return False - await _async_migrate_data(hass, entry, protect) + await async_migrate_data(hass, entry, protect) if entry.unique_id is None: hass.config_entries.async_update_entry(entry, unique_id=nvr_info.mac) diff --git a/homeassistant/components/unifiprotect/migrate.py b/homeassistant/components/unifiprotect/migrate.py new file mode 100644 index 00000000000..dcba0b504c9 --- /dev/null +++ b/homeassistant/components/unifiprotect/migrate.py @@ -0,0 +1,83 @@ +"""UniFi Protect data migrations.""" +from __future__ import annotations + +import logging + +from pyunifiprotect import ProtectApiClient + +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import Platform +from homeassistant.core import HomeAssistant +from homeassistant.helpers import entity_registry as er + +from .const import DEVICES_THAT_ADOPT + +_LOGGER = logging.getLogger(__name__) + + +async def async_migrate_data( + hass: HomeAssistant, entry: ConfigEntry, protect: ProtectApiClient +) -> None: + """Run all valid UniFi Protect data migrations.""" + + _LOGGER.debug("Start Migrate: async_migrate_buttons") + await async_migrate_buttons(hass, entry, protect) + _LOGGER.debug("Completed Migrate: async_migrate_buttons") + + +async def async_migrate_buttons( + hass: HomeAssistant, entry: ConfigEntry, protect: ProtectApiClient +) -> None: + """ + Migrate existing Reboot button unique IDs from {device_id} to {deivce_id}_reboot. + + This allows for additional types of buttons that are outside of just a reboot button. + + Added in 2022.6.0. + """ + + registry = er.async_get(hass) + to_migrate = [] + for entity in er.async_entries_for_config_entry(registry, entry.entry_id): + if entity.domain == Platform.BUTTON and "_" not in entity.unique_id: + _LOGGER.debug("Button %s needs migration", entity.entity_id) + to_migrate.append(entity) + + if len(to_migrate) == 0: + _LOGGER.debug("No button entities need migration") + return + + bootstrap = await protect.get_bootstrap() + count = 0 + for button in to_migrate: + device = None + for model in DEVICES_THAT_ADOPT: + attr = f"{model.value}s" + device = getattr(bootstrap, attr).get(button.unique_id) + if device is not None: + break + + if device is None: + continue + + new_unique_id = f"{device.id}_reboot" + _LOGGER.debug( + "Migrating entity %s (old unique_id: %s, new unique_id: %s)", + button.entity_id, + button.unique_id, + new_unique_id, + ) + try: + registry.async_update_entity(button.entity_id, new_unique_id=new_unique_id) + except ValueError: + _LOGGER.warning( + "Could not migrate entity %s (old unique_id: %s, new unique_id: %s)", + button.entity_id, + button.unique_id, + new_unique_id, + ) + else: + count += 1 + + if count < len(to_migrate): + _LOGGER.warning("Failed to migate %s reboot buttons", len(to_migrate) - count) diff --git a/tests/components/unifiprotect/test_init.py b/tests/components/unifiprotect/test_init.py index cf899d854fd..68f171b52bf 100644 --- a/tests/components/unifiprotect/test_init.py +++ b/tests/components/unifiprotect/test_init.py @@ -11,7 +11,6 @@ from pyunifiprotect.data import NVR, Light from homeassistant.components.unifiprotect.const import CONF_DISABLE_RTSP, DOMAIN from homeassistant.config_entries import ConfigEntry, ConfigEntryState -from homeassistant.const import Platform from homeassistant.core import HomeAssistant from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.setup import async_setup_component @@ -199,149 +198,6 @@ async def test_setup_starts_discovery( assert len(hass.config_entries.flow.async_progress_by_handler(DOMAIN)) == 1 -async def test_migrate_reboot_button( - hass: HomeAssistant, mock_entry: MockEntityFixture, mock_light: Light -): - """Test migrating unique ID of reboot button.""" - - light1 = mock_light.copy() - light1._api = mock_entry.api - light1.name = "Test Light 1" - light1.id = "lightid1" - - light2 = mock_light.copy() - light2._api = mock_entry.api - light2.name = "Test Light 2" - light2.id = "lightid2" - mock_entry.api.bootstrap.lights = { - light1.id: light1, - light2.id: light2, - } - mock_entry.api.get_bootstrap = AsyncMock(return_value=mock_entry.api.bootstrap) - - registry = er.async_get(hass) - registry.async_get_or_create( - Platform.BUTTON, DOMAIN, light1.id, config_entry=mock_entry.entry - ) - registry.async_get_or_create( - Platform.BUTTON, - DOMAIN, - f"{light2.id}_reboot", - config_entry=mock_entry.entry, - ) - - await hass.config_entries.async_setup(mock_entry.entry.entry_id) - await hass.async_block_till_done() - - assert mock_entry.entry.state == ConfigEntryState.LOADED - assert mock_entry.api.update.called - assert mock_entry.entry.unique_id == mock_entry.api.bootstrap.nvr.mac - - buttons = [] - for entity in er.async_entries_for_config_entry( - registry, mock_entry.entry.entry_id - ): - if entity.domain == Platform.BUTTON.value: - buttons.append(entity) - print(entity.entity_id) - assert len(buttons) == 2 - - assert registry.async_get(f"{Platform.BUTTON}.test_light_1_reboot_device") is None - assert registry.async_get(f"{Platform.BUTTON}.test_light_1_reboot_device_2") is None - light = registry.async_get(f"{Platform.BUTTON}.unifiprotect_lightid1") - assert light is not None - assert light.unique_id == f"{light1.id}_reboot" - - assert registry.async_get(f"{Platform.BUTTON}.test_light_2_reboot_device") is None - assert registry.async_get(f"{Platform.BUTTON}.test_light_2_reboot_device_2") is None - light = registry.async_get(f"{Platform.BUTTON}.unifiprotect_lightid2_reboot") - assert light is not None - assert light.unique_id == f"{light2.id}_reboot" - - -async def test_migrate_reboot_button_no_device( - hass: HomeAssistant, mock_entry: MockEntityFixture, mock_light: Light -): - """Test migrating unique ID of reboot button if UniFi Protect device ID changed.""" - - light1 = mock_light.copy() - light1._api = mock_entry.api - light1.name = "Test Light 1" - light1.id = "lightid1" - - mock_entry.api.bootstrap.lights = { - light1.id: light1, - } - mock_entry.api.get_bootstrap = AsyncMock(return_value=mock_entry.api.bootstrap) - - registry = er.async_get(hass) - registry.async_get_or_create( - Platform.BUTTON, DOMAIN, "lightid2", config_entry=mock_entry.entry - ) - - await hass.config_entries.async_setup(mock_entry.entry.entry_id) - await hass.async_block_till_done() - - assert mock_entry.entry.state == ConfigEntryState.LOADED - assert mock_entry.api.update.called - assert mock_entry.entry.unique_id == mock_entry.api.bootstrap.nvr.mac - - buttons = [] - for entity in er.async_entries_for_config_entry( - registry, mock_entry.entry.entry_id - ): - if entity.domain == Platform.BUTTON.value: - buttons.append(entity) - assert len(buttons) == 2 - - light = registry.async_get(f"{Platform.BUTTON}.unifiprotect_lightid2") - assert light is not None - assert light.unique_id == "lightid2" - - -async def test_migrate_reboot_button_fail( - hass: HomeAssistant, mock_entry: MockEntityFixture, mock_light: Light -): - """Test migrating unique ID of reboot button.""" - - light1 = mock_light.copy() - light1._api = mock_entry.api - light1.name = "Test Light 1" - light1.id = "lightid1" - - mock_entry.api.bootstrap.lights = { - light1.id: light1, - } - mock_entry.api.get_bootstrap = AsyncMock(return_value=mock_entry.api.bootstrap) - - registry = er.async_get(hass) - registry.async_get_or_create( - Platform.BUTTON, - DOMAIN, - light1.id, - config_entry=mock_entry.entry, - suggested_object_id=light1.name, - ) - registry.async_get_or_create( - Platform.BUTTON, - DOMAIN, - f"{light1.id}_reboot", - config_entry=mock_entry.entry, - suggested_object_id=light1.name, - ) - - await hass.config_entries.async_setup(mock_entry.entry.entry_id) - await hass.async_block_till_done() - - assert mock_entry.entry.state == ConfigEntryState.LOADED - assert mock_entry.api.update.called - assert mock_entry.entry.unique_id == mock_entry.api.bootstrap.nvr.mac - - light = registry.async_get(f"{Platform.BUTTON}.test_light_1") - assert light is not None - assert light.unique_id == f"{light1.id}" - - async def test_device_remove_devices( hass: HomeAssistant, mock_entry: MockEntityFixture, diff --git a/tests/components/unifiprotect/test_migrate.py b/tests/components/unifiprotect/test_migrate.py new file mode 100644 index 00000000000..756672bcbca --- /dev/null +++ b/tests/components/unifiprotect/test_migrate.py @@ -0,0 +1,158 @@ +"""Test the UniFi Protect setup flow.""" +# pylint: disable=protected-access +from __future__ import annotations + +from unittest.mock import AsyncMock + +from pyunifiprotect.data import Light + +from homeassistant.components.unifiprotect.const import DOMAIN +from homeassistant.config_entries import ConfigEntryState +from homeassistant.const import Platform +from homeassistant.core import HomeAssistant +from homeassistant.helpers import entity_registry as er + +from .conftest import MockEntityFixture + + +async def test_migrate_reboot_button( + hass: HomeAssistant, mock_entry: MockEntityFixture, mock_light: Light +): + """Test migrating unique ID of reboot button.""" + + light1 = mock_light.copy() + light1._api = mock_entry.api + light1.name = "Test Light 1" + light1.id = "lightid1" + + light2 = mock_light.copy() + light2._api = mock_entry.api + light2.name = "Test Light 2" + light2.id = "lightid2" + mock_entry.api.bootstrap.lights = { + light1.id: light1, + light2.id: light2, + } + mock_entry.api.get_bootstrap = AsyncMock(return_value=mock_entry.api.bootstrap) + + registry = er.async_get(hass) + registry.async_get_or_create( + Platform.BUTTON, DOMAIN, light1.id, config_entry=mock_entry.entry + ) + registry.async_get_or_create( + Platform.BUTTON, + DOMAIN, + f"{light2.id}_reboot", + config_entry=mock_entry.entry, + ) + + await hass.config_entries.async_setup(mock_entry.entry.entry_id) + await hass.async_block_till_done() + + assert mock_entry.entry.state == ConfigEntryState.LOADED + assert mock_entry.api.update.called + assert mock_entry.entry.unique_id == mock_entry.api.bootstrap.nvr.mac + + buttons = [] + for entity in er.async_entries_for_config_entry( + registry, mock_entry.entry.entry_id + ): + if entity.domain == Platform.BUTTON.value: + buttons.append(entity) + print(entity.entity_id) + assert len(buttons) == 2 + + assert registry.async_get(f"{Platform.BUTTON}.test_light_1_reboot_device") is None + assert registry.async_get(f"{Platform.BUTTON}.test_light_1_reboot_device_2") is None + light = registry.async_get(f"{Platform.BUTTON}.unifiprotect_lightid1") + assert light is not None + assert light.unique_id == f"{light1.id}_reboot" + + assert registry.async_get(f"{Platform.BUTTON}.test_light_2_reboot_device") is None + assert registry.async_get(f"{Platform.BUTTON}.test_light_2_reboot_device_2") is None + light = registry.async_get(f"{Platform.BUTTON}.unifiprotect_lightid2_reboot") + assert light is not None + assert light.unique_id == f"{light2.id}_reboot" + + +async def test_migrate_reboot_button_no_device( + hass: HomeAssistant, mock_entry: MockEntityFixture, mock_light: Light +): + """Test migrating unique ID of reboot button if UniFi Protect device ID changed.""" + + light1 = mock_light.copy() + light1._api = mock_entry.api + light1.name = "Test Light 1" + light1.id = "lightid1" + + mock_entry.api.bootstrap.lights = { + light1.id: light1, + } + mock_entry.api.get_bootstrap = AsyncMock(return_value=mock_entry.api.bootstrap) + + registry = er.async_get(hass) + registry.async_get_or_create( + Platform.BUTTON, DOMAIN, "lightid2", config_entry=mock_entry.entry + ) + + await hass.config_entries.async_setup(mock_entry.entry.entry_id) + await hass.async_block_till_done() + + assert mock_entry.entry.state == ConfigEntryState.LOADED + assert mock_entry.api.update.called + assert mock_entry.entry.unique_id == mock_entry.api.bootstrap.nvr.mac + + buttons = [] + for entity in er.async_entries_for_config_entry( + registry, mock_entry.entry.entry_id + ): + if entity.domain == Platform.BUTTON.value: + buttons.append(entity) + assert len(buttons) == 2 + + light = registry.async_get(f"{Platform.BUTTON}.unifiprotect_lightid2") + assert light is not None + assert light.unique_id == "lightid2" + + +async def test_migrate_reboot_button_fail( + hass: HomeAssistant, mock_entry: MockEntityFixture, mock_light: Light +): + """Test migrating unique ID of reboot button.""" + + light1 = mock_light.copy() + light1._api = mock_entry.api + light1.name = "Test Light 1" + light1.id = "lightid1" + + mock_entry.api.bootstrap.lights = { + light1.id: light1, + } + mock_entry.api.get_bootstrap = AsyncMock(return_value=mock_entry.api.bootstrap) + + registry = er.async_get(hass) + registry.async_get_or_create( + Platform.BUTTON, + DOMAIN, + light1.id, + config_entry=mock_entry.entry, + suggested_object_id=light1.name, + ) + registry.async_get_or_create( + Platform.BUTTON, + DOMAIN, + f"{light1.id}_reboot", + config_entry=mock_entry.entry, + suggested_object_id=light1.name, + ) + + await hass.config_entries.async_setup(mock_entry.entry.entry_id) + await hass.async_block_till_done() + + assert mock_entry.entry.state == ConfigEntryState.LOADED + assert mock_entry.api.update.called + assert mock_entry.entry.unique_id == mock_entry.api.bootstrap.nvr.mac + + light = registry.async_get(f"{Platform.BUTTON}.test_light_1") + assert light is not None + assert light.unique_id == f"{light1.id}"