Move update service registration to services module (#182826)

homeassistant/components/update/__init__.py

@@ -17,14 +17,13 @@ from homeassistant.const import (
     EntityCategory,
     EntityStateAttribute,
 )
-from homeassistant.core import HomeAssistant, ServiceCall
+from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.entity import ABCCachedProperties, EntityDescription
 from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import ConfigType
-from homeassistant.util.hass_dict import HassKey
 
 from .const import (  # noqa: F401
     ATTR_AUTO_UPDATE,
@@ -39,6 +38,7 @@ from .const import (  # noqa: F401
     ATTR_TITLE,
     ATTR_UPDATE_PERCENTAGE,
     ATTR_VERSION,
+    DATA_COMPONENT,
     DEVICE_CLASSES_SCHEMA,
     DOMAIN,
     SERVICE_INSTALL,
@@ -47,10 +47,10 @@ from .const import (  # noqa: F401
     UpdateEntityFeature,
     UpdateEntityStateAttribute,
 )
+from .services import async_setup_services
 
 _LOGGER = logging.getLogger(__name__)
 
-DATA_COMPONENT: HassKey[EntityComponent[UpdateEntity]] = HassKey(DOMAIN)
 ENTITY_ID_FORMAT: Final = DOMAIN + ".{}"
 PLATFORM_SCHEMA = cv.PLATFORM_SCHEMA
 PLATFORM_SCHEMA_BASE = cv.PLATFORM_SCHEMA_BASE
@@ -85,29 +85,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     )
     await component.async_setup(config)
 
-    component.async_register_entity_service(
-        SERVICE_INSTALL,
-        {
-            probatio.Optional(ATTR_VERSION): cv.string,
-            probatio.Optional(ATTR_BACKUP, default=False): cv.boolean,
-        },
-        async_install,
-        [UpdateEntityFeature.INSTALL],
-        admin_only=True,
-    )
-
-    component.async_register_entity_service(
-        SERVICE_SKIP,
-        None,
-        async_skip,
-        admin_only=True,
-    )
-    component.async_register_entity_service(
-        "clear_skipped",
-        None,
-        async_clear_skipped,
-        admin_only=True,
-    )
+    async_setup_services(hass)
 
     websocket_api.async_register_command(hass, websocket_release_notes)
 
@@ -124,73 +102,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return await hass.data[DATA_COMPONENT].async_unload_entry(entry)
 
 
-async def async_install(entity: UpdateEntity, service_call: ServiceCall) -> None:
-    """Service call wrapper to validate the call."""
-    # If version is not specified, but no update is available.
-    if (version := service_call.data.get(ATTR_VERSION)) is None and (
-        entity.installed_version == entity.latest_version
-        or entity.latest_version is None
-    ):
-        raise HomeAssistantError(
-            translation_domain=DOMAIN,
-            translation_key="no_update_available",
-            translation_placeholders={"entity_id": entity.entity_id},
-        )
-
-    # If version is specified, but not supported by the entity.
-    if (
-        version is not None
-        and UpdateEntityFeature.SPECIFIC_VERSION not in entity.supported_features
-    ):
-        raise HomeAssistantError(
-            translation_domain=DOMAIN,
-            translation_key="specific_version_not_supported",
-            translation_placeholders={"entity_id": entity.entity_id},
-        )
-
-    # If backup is requested, but not supported by the entity.
-    if (
-        backup := service_call.data[ATTR_BACKUP]
-    ) and UpdateEntityFeature.BACKUP not in entity.supported_features:
-        raise HomeAssistantError(
-            translation_domain=DOMAIN,
-            translation_key="backup_not_supported",
-            translation_placeholders={"entity_id": entity.entity_id},
-        )
-
-    # Update is already in progress.
-    if entity.in_progress is not False:
-        raise HomeAssistantError(
-            translation_domain=DOMAIN,
-            translation_key="update_in_progress",
-            translation_placeholders={"entity_id": entity.entity_id},
-        )
-
-    await entity.async_install_with_progress(version, backup)
-
-
-async def async_skip(entity: UpdateEntity, service_call: ServiceCall) -> None:
-    """Service call wrapper to validate the call."""
-    if entity.auto_update:
-        raise HomeAssistantError(
-            translation_domain=DOMAIN,
-            translation_key="skip_not_supported",
-            translation_placeholders={"entity_id": entity.entity_id},
-        )
-    await entity.async_skip()
-
-
-async def async_clear_skipped(entity: UpdateEntity, service_call: ServiceCall) -> None:
-    """Service call wrapper to validate the call."""
-    if entity.auto_update:
-        raise HomeAssistantError(
-            translation_domain=DOMAIN,
-            translation_key="clear_skipped_not_supported",
-            translation_placeholders={"entity_id": entity.entity_id},
-        )
-    await entity.async_clear_skipped()
-
-
 class UpdateEntityDescription(EntityDescription, frozen_or_thawed=True):
     """A class that describes update entities."""
 

homeassistant/components/update/const.py

@@ -1,11 +1,19 @@
 """Constants for the update component."""
 
 from enum import IntFlag, StrEnum
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 import probatio
 
+from homeassistant.util.hass_dict import HassKey
+
+if TYPE_CHECKING:
+    from homeassistant.helpers.entity_component import EntityComponent
+
+    from . import UpdateEntity
+
 DOMAIN: Final = "update"
+DATA_COMPONENT: HassKey[EntityComponent[UpdateEntity]] = HassKey(DOMAIN)
 
 
 class UpdateEntityStateAttribute(StrEnum):

homeassistant/components/update/services.py

@@ -0,0 +1,119 @@
+"""Services for the Update integration."""
+
+from typing import TYPE_CHECKING
+
+import probatio
+
+from homeassistant.core import HomeAssistant, ServiceCall, callback
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import config_validation as cv
+
+from .const import (
+    ATTR_BACKUP,
+    ATTR_VERSION,
+    DATA_COMPONENT,
+    DOMAIN,
+    SERVICE_INSTALL,
+    SERVICE_SKIP,
+    UpdateEntityFeature,
+)
+
+if TYPE_CHECKING:
+    from . import UpdateEntity
+
+
+async def _async_install(entity: UpdateEntity, service_call: ServiceCall) -> None:
+    """Service call wrapper to validate the call."""
+    # If version is not specified, but no update is available.
+    if (version := service_call.data.get(ATTR_VERSION)) is None and (
+        entity.installed_version == entity.latest_version
+        or entity.latest_version is None
+    ):
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="no_update_available",
+            translation_placeholders={"entity_id": entity.entity_id},
+        )
+
+    # If version is specified, but not supported by the entity.
+    if (
+        version is not None
+        and UpdateEntityFeature.SPECIFIC_VERSION not in entity.supported_features
+    ):
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="specific_version_not_supported",
+            translation_placeholders={"entity_id": entity.entity_id},
+        )
+
+    # If backup is requested, but not supported by the entity.
+    if (
+        backup := service_call.data[ATTR_BACKUP]
+    ) and UpdateEntityFeature.BACKUP not in entity.supported_features:
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="backup_not_supported",
+            translation_placeholders={"entity_id": entity.entity_id},
+        )
+
+    # Update is already in progress.
+    if entity.in_progress is not False:
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="update_in_progress",
+            translation_placeholders={"entity_id": entity.entity_id},
+        )
+
+    await entity.async_install_with_progress(version, backup)
+
+
+async def _async_skip(entity: UpdateEntity, service_call: ServiceCall) -> None:
+    """Service call wrapper to validate the call."""
+    if entity.auto_update:
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="skip_not_supported",
+            translation_placeholders={"entity_id": entity.entity_id},
+        )
+    await entity.async_skip()
+
+
+async def _async_clear_skipped(entity: UpdateEntity, service_call: ServiceCall) -> None:
+    """Service call wrapper to validate the call."""
+    if entity.auto_update:
+        raise HomeAssistantError(
+            translation_domain=DOMAIN,
+            translation_key="clear_skipped_not_supported",
+            translation_placeholders={"entity_id": entity.entity_id},
+        )
+    await entity.async_clear_skipped()
+
+
+@callback
+def async_setup_services(hass: HomeAssistant) -> None:
+    """Register the update services."""
+    component = hass.data[DATA_COMPONENT]
+
+    component.async_register_entity_service(
+        SERVICE_INSTALL,
+        {
+            probatio.Optional(ATTR_VERSION): cv.string,
+            probatio.Optional(ATTR_BACKUP, default=False): cv.boolean,
+        },
+        _async_install,
+        [UpdateEntityFeature.INSTALL],
+        admin_only=True,
+    )
+
+    component.async_register_entity_service(
+        SERVICE_SKIP,
+        None,
+        _async_skip,
+        admin_only=True,
+    )
+    component.async_register_entity_service(
+        "clear_skipped",
+        None,
+        _async_clear_skipped,
+        admin_only=True,
+    )