Simplify UniFi Protect service setup/cleanup (#63908)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/unifiprotect/__init__.py

@@ -3,14 +3,13 @@ from __future__ import annotations
 
 import asyncio
 from datetime import timedelta
-import functools
 import logging
 
 from aiohttp import CookieJar
 from aiohttp.client_exceptions import ServerDisconnectedError
 from pyunifiprotect import NotAuthorized, NvrError, ProtectApiClient
 
-from homeassistant.config_entries import ConfigEntry, ConfigEntryState
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -24,22 +23,17 @@ from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
 
 from .const import (
-    ALL_GLOBAL_SERIVCES,
     CONF_ALL_UPDATES,
     CONF_OVERRIDE_CHOST,
     DEFAULT_SCAN_INTERVAL,
     DEVICES_FOR_SUBSCRIBE,
     DOMAIN,
-    DOORBELL_TEXT_SCHEMA,
     MIN_REQUIRED_PROTECT_V,
     OUTDATED_LOG_MESSAGE,
     PLATFORMS,
-    SERVICE_ADD_DOORBELL_TEXT,
-    SERVICE_REMOVE_DOORBELL_TEXT,
-    SERVICE_SET_DEFAULT_DOORBELL_TEXT,
 )
 from .data import ProtectData
-from .services import add_doorbell_text, remove_doorbell_text, set_default_doorbell_text
+from .services import async_cleanup_services, async_setup_services
 from .views import ThumbnailProxyView
 
 _LOGGER = logging.getLogger(__name__)
@@ -89,29 +83,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = data_service
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
-
+    async_setup_services(hass)
-    services = [
-        (
-            SERVICE_ADD_DOORBELL_TEXT,
-            functools.partial(add_doorbell_text, hass),
-            DOORBELL_TEXT_SCHEMA,
-        ),
-        (
-            SERVICE_REMOVE_DOORBELL_TEXT,
-            functools.partial(remove_doorbell_text, hass),
-            DOORBELL_TEXT_SCHEMA,
-        ),
-        (
-            SERVICE_SET_DEFAULT_DOORBELL_TEXT,
-            functools.partial(set_default_doorbell_text, hass),
-            DOORBELL_TEXT_SCHEMA,
-        ),
-    ]
-    for name, method, schema in services:
-        if hass.services.has_service(DOMAIN, name):
-            continue
-        hass.services.async_register(DOMAIN, name, method, schema=schema)
-
     hass.http.register_view(ThumbnailProxyView(hass))
 
     entry.async_on_unload(entry.add_update_listener(_async_options_updated))
@@ -133,14 +105,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         data: ProtectData = hass.data[DOMAIN][entry.entry_id]
         await data.async_stop()
         hass.data[DOMAIN].pop(entry.entry_id)
-
+        async_cleanup_services(hass)
-    loaded_entries = [
-        entry
-        for entry in hass.config_entries.async_entries(DOMAIN)
-        if entry.state == ConfigEntryState.LOADED
-    ]
-    if len(loaded_entries) == 1:
-        for name in ALL_GLOBAL_SERIVCES:
-            hass.services.async_remove(DOMAIN, name)
 
     return bool(unload_ok)

homeassistant/components/unifiprotect/const.py

@@ -1,10 +1,8 @@
 """Constant definitions for UniFi Protect Integration."""
 
 from pyunifiprotect.data.types import ModelType, Version
-import voluptuous as vol
 
-from homeassistant.const import ATTR_DEVICE_ID, ATTR_ENTITY_ID, Platform
+from homeassistant.const import Platform
-from homeassistant.helpers import config_validation as cv
 
 DOMAIN = "unifiprotect"
 
@@ -49,16 +47,6 @@ OUTDATED_LOG_MESSAGE = "You are running v%s of UniFi Protect. Minimum required v
 
 TYPE_EMPTY_VALUE = ""
 
-SERVICE_ADD_DOORBELL_TEXT = "add_doorbell_text"
-SERVICE_REMOVE_DOORBELL_TEXT = "remove_doorbell_text"
-SERVICE_SET_DEFAULT_DOORBELL_TEXT = "set_default_doorbell_text"
-
-ALL_GLOBAL_SERIVCES = [
-    SERVICE_ADD_DOORBELL_TEXT,
-    SERVICE_REMOVE_DOORBELL_TEXT,
-    SERVICE_SET_DEFAULT_DOORBELL_TEXT,
-]
-
 PLATFORMS = [
     Platform.BINARY_SENSOR,
     Platform.BUTTON,
@@ -70,43 +58,3 @@ PLATFORMS = [
     Platform.SENSOR,
     Platform.SWITCH,
 ]
-
-
-DOORBELL_TEXT_SCHEMA = vol.All(
-    vol.Schema(
-        {
-            **cv.ENTITY_SERVICE_FIELDS,
-            vol.Required(ATTR_MESSAGE): cv.string,
-        },
-    ),
-    cv.has_at_least_one_key(ATTR_DEVICE_ID),
-)
-
-GENERATE_DATA_SCHEMA = vol.All(
-    vol.Schema(
-        {
-            **cv.ENTITY_SERVICE_FIELDS,
-            vol.Required(ATTR_DURATION): vol.Coerce(int),
-            vol.Required(ATTR_ANONYMIZE): vol.Coerce(bool),
-        },
-    ),
-    cv.has_at_least_one_key(ATTR_DEVICE_ID),
-)
-
-PROFILE_WS_SCHEMA = vol.All(
-    vol.Schema(
-        {
-            **cv.ENTITY_SERVICE_FIELDS,
-            vol.Required(ATTR_DURATION): vol.Coerce(int),
-        },
-    ),
-    cv.has_at_least_one_key(ATTR_DEVICE_ID),
-)
-
-SET_DOORBELL_LCD_MESSAGE_SCHEMA = vol.Schema(
-    {
-        vol.Required(ATTR_ENTITY_ID): cv.entity_ids,
-        vol.Required(ATTR_MESSAGE): cv.string,
-        vol.Optional(ATTR_DURATION, default="None"): cv.string,
-    }
-)

homeassistant/components/unifiprotect/services.py

@@ -2,20 +2,44 @@
 from __future__ import annotations
 
 import asyncio
+import functools
 from typing import Any
 
 from pydantic import ValidationError
 from pyunifiprotect.api import ProtectApiClient
 from pyunifiprotect.exceptions import BadRequest
+import voluptuous as vol
 
+from homeassistant.config_entries import ConfigEntryState
+from homeassistant.const import ATTR_DEVICE_ID
 from homeassistant.core import HomeAssistant, ServiceCall, callback
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers import config_validation as cv, device_registry as dr
 from homeassistant.helpers.service import async_extract_referenced_entity_ids
 
 from .const import ATTR_MESSAGE, DOMAIN
 from .data import ProtectData
 
+SERVICE_ADD_DOORBELL_TEXT = "add_doorbell_text"
+SERVICE_REMOVE_DOORBELL_TEXT = "remove_doorbell_text"
+SERVICE_SET_DEFAULT_DOORBELL_TEXT = "set_default_doorbell_text"
+
+ALL_GLOBAL_SERIVCES = [
+    SERVICE_ADD_DOORBELL_TEXT,
+    SERVICE_REMOVE_DOORBELL_TEXT,
+    SERVICE_SET_DEFAULT_DOORBELL_TEXT,
+]
+
+DOORBELL_TEXT_SCHEMA = vol.All(
+    vol.Schema(
+        {
+            **cv.ENTITY_SERVICE_FIELDS,
+            vol.Required(ATTR_MESSAGE): cv.string,
+        },
+    ),
+    cv.has_at_least_one_key(ATTR_DEVICE_ID),
+)
+
 
 def _async_all_ufp_instances(hass: HomeAssistant) -> list[ProtectApiClient]:
     """All active UFP instances."""
@@ -110,3 +134,40 @@ async def set_default_doorbell_text(hass: HomeAssistant, call: ServiceCall) -> N
     message: str = call.data[ATTR_MESSAGE]
     instances = _async_get_protect_from_call(hass, call)
     await _async_call_nvr(instances, "set_default_doorbell_message", message)
+
+
+def async_setup_services(hass: HomeAssistant) -> None:
+    """Set up the global UniFi Protect services."""
+    services = [
+        (
+            SERVICE_ADD_DOORBELL_TEXT,
+            functools.partial(add_doorbell_text, hass),
+            DOORBELL_TEXT_SCHEMA,
+        ),
+        (
+            SERVICE_REMOVE_DOORBELL_TEXT,
+            functools.partial(remove_doorbell_text, hass),
+            DOORBELL_TEXT_SCHEMA,
+        ),
+        (
+            SERVICE_SET_DEFAULT_DOORBELL_TEXT,
+            functools.partial(set_default_doorbell_text, hass),
+            DOORBELL_TEXT_SCHEMA,
+        ),
+    ]
+    for name, method, schema in services:
+        if hass.services.has_service(DOMAIN, name):
+            continue
+        hass.services.async_register(DOMAIN, name, method, schema=schema)
+
+
+def async_cleanup_services(hass: HomeAssistant) -> None:
+    """Cleanup global UniFi Protect services (if all config entries unloaded)."""
+    loaded_entries = [
+        entry
+        for entry in hass.config_entries.async_entries(DOMAIN)
+        if entry.state == ConfigEntryState.LOADED
+    ]
+    if len(loaded_entries) == 1:
+        for name in ALL_GLOBAL_SERIVCES:
+            hass.services.async_remove(DOMAIN, name)

tests/components/unifiprotect/test_services.py

@@ -8,9 +8,8 @@ import pytest
 from pyunifiprotect.data import Light
 from pyunifiprotect.exceptions import BadRequest
 
-from homeassistant.components.unifiprotect.const import (
+from homeassistant.components.unifiprotect.const import ATTR_MESSAGE, DOMAIN
-    ATTR_MESSAGE,
+from homeassistant.components.unifiprotect.services import (
-    DOMAIN,
     SERVICE_ADD_DOORBELL_TEXT,
     SERVICE_REMOVE_DOORBELL_TEXT,
     SERVICE_SET_DEFAULT_DOORBELL_TEXT,