diff --git a/homeassistant/components/unifiprotect/__init__.py b/homeassistant/components/unifiprotect/__init__.py index 942c533b59e..c4a6bc88068 100644 --- a/homeassistant/components/unifiprotect/__init__.py +++ b/homeassistant/components/unifiprotect/__init__.py @@ -5,6 +5,7 @@ from datetime import timedelta import logging from aiohttp.client_exceptions import ServerDisconnectedError +from pyunifiprotect.data import Bootstrap from pyunifiprotect.exceptions import ClientError, NotAuthorized # Import the test_util.anonymize module from the pyunifiprotect package @@ -128,7 +129,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ) try: - await _async_setup_entry(hass, entry, data_service) + await _async_setup_entry(hass, entry, data_service, bootstrap) except Exception as err: if await nvr_info.get_is_prerelease(): # If they are running a pre-release, its quite common for setup @@ -156,9 +157,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def _async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, data_service: ProtectData + hass: HomeAssistant, + entry: ConfigEntry, + data_service: ProtectData, + bootstrap: Bootstrap, ) -> None: - await async_migrate_data(hass, entry, data_service.api) + await async_migrate_data(hass, entry, data_service.api, bootstrap) await data_service.async_setup() if not data_service.last_update_success: diff --git a/homeassistant/components/unifiprotect/migrate.py b/homeassistant/components/unifiprotect/migrate.py index db1e82d9914..3a6dde653b4 100644 --- a/homeassistant/components/unifiprotect/migrate.py +++ b/homeassistant/components/unifiprotect/migrate.py @@ -3,47 +3,39 @@ from __future__ import annotations import logging -from aiohttp.client_exceptions import ServerDisconnectedError from pyunifiprotect import ProtectApiClient from pyunifiprotect.data import NVR, Bootstrap, ProtectAdoptableDeviceModel -from pyunifiprotect.exceptions import ClientError from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers import entity_registry as er _LOGGER = logging.getLogger(__name__) async def async_migrate_data( - hass: HomeAssistant, entry: ConfigEntry, protect: ProtectApiClient + hass: HomeAssistant, + entry: ConfigEntry, + protect: ProtectApiClient, + bootstrap: Bootstrap, ) -> None: """Run all valid UniFi Protect data migrations.""" _LOGGER.debug("Start Migrate: async_migrate_buttons") - await async_migrate_buttons(hass, entry, protect) + await async_migrate_buttons(hass, entry, protect, bootstrap) _LOGGER.debug("Completed Migrate: async_migrate_buttons") _LOGGER.debug("Start Migrate: async_migrate_device_ids") - await async_migrate_device_ids(hass, entry, protect) + await async_migrate_device_ids(hass, entry, protect, bootstrap) _LOGGER.debug("Completed Migrate: async_migrate_device_ids") -async def async_get_bootstrap(protect: ProtectApiClient) -> Bootstrap: - """Get UniFi Protect bootstrap or raise appropriate HA error.""" - - try: - bootstrap = await protect.get_bootstrap() - except (TimeoutError, ClientError, ServerDisconnectedError) as err: - raise ConfigEntryNotReady from err - - return bootstrap - - async def async_migrate_buttons( - hass: HomeAssistant, entry: ConfigEntry, protect: ProtectApiClient + hass: HomeAssistant, + entry: ConfigEntry, + protect: ProtectApiClient, + bootstrap: Bootstrap, ) -> None: """Migrate existing Reboot button unique IDs from {device_id} to {deivce_id}_reboot. @@ -63,7 +55,6 @@ async def async_migrate_buttons( _LOGGER.debug("No button entities need migration") return - bootstrap = await async_get_bootstrap(protect) count = 0 for button in to_migrate: device = bootstrap.get_device_from_id(button.unique_id) @@ -94,7 +85,10 @@ async def async_migrate_buttons( async def async_migrate_device_ids( - hass: HomeAssistant, entry: ConfigEntry, protect: ProtectApiClient + hass: HomeAssistant, + entry: ConfigEntry, + protect: ProtectApiClient, + bootstrap: Bootstrap, ) -> None: """Migrate unique IDs from {device_id}_{name} format to {mac}_{name} format. @@ -119,7 +113,6 @@ async def async_migrate_device_ids( _LOGGER.debug("No entities need migration to MAC address ID") return - bootstrap = await async_get_bootstrap(protect) count = 0 for entity in to_migrate: parts = entity.unique_id.split("_")