evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor migration code for UniFi Protect (#73499)

Browse source
This commit is contained in:
Christopher Bailey 2022-06-18 13:56:28 -04:00 • committed by GitHub
parent f4c3bd7e00
commit 691d49f23b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 244 additions and 201 deletions
Download patch file Download diff file Expand all files Collapse all files

60
homeassistant/components/unifiprotect/__init__.py
View file

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

83
homeassistant/components/unifiprotect/migrate.py Normal file
View file

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

144
tests/components/unifiprotect/test_init.py
View file

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

158
tests/components/unifiprotect/test_migrate.py Normal file
View file

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