Refactor migration code for UniFi Protect (#73499)

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)
 

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)

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,

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