Move Guardian services to entity platform services (#37189)

homeassistant/components/guardian/__init__.py

@@ -3,34 +3,17 @@ import asyncio
 from datetime import timedelta
 
 from aioguardian import Client
-from aioguardian.commands.system import (
-    DEFAULT_FIRMWARE_UPGRADE_FILENAME,
-    DEFAULT_FIRMWARE_UPGRADE_PORT,
-    DEFAULT_FIRMWARE_UPGRADE_URL,
-)
 from aioguardian.errors import GuardianError
-import voluptuous as vol
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import (
+from homeassistant.const import ATTR_ATTRIBUTION, CONF_IP_ADDRESS
-    ATTR_ATTRIBUTION,
-    CONF_FILENAME,
-    CONF_IP_ADDRESS,
-    CONF_PORT,
-    CONF_URL,
-)
 from homeassistant.core import HomeAssistant, callback
-import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import (
     async_dispatcher_connect,
     async_dispatcher_send,
 )
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.event import async_track_time_interval
-from homeassistant.helpers.service import (
-    async_register_admin_service,
-    verify_domain_control,
-)
 
 from .const import (
     CONF_UID,
@@ -61,16 +44,6 @@ DEFAULT_SCAN_INTERVAL = timedelta(seconds=30)
 
 PLATFORMS = ["binary_sensor", "sensor", "switch"]
 
-SERVICE_UPGRADE_FIRMWARE_SCHEMA = vol.Schema(
-    {
-        vol.Optional(CONF_URL, default=DEFAULT_FIRMWARE_UPGRADE_URL): cv.url,
-        vol.Optional(CONF_PORT, default=DEFAULT_FIRMWARE_UPGRADE_PORT): cv.port,
-        vol.Optional(
-            CONF_FILENAME, default=DEFAULT_FIRMWARE_UPGRADE_FILENAME
-        ): cv.string,
-    }
-)
-
 
 @callback
 def async_get_api_category(entity_kind: str):
@@ -86,8 +59,6 @@ async def async_setup(hass: HomeAssistant, config: dict):
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     """Set up Elexa Guardian from a config entry."""
-    _verify_domain_control = verify_domain_control(hass, DOMAIN)
-
     guardian = Guardian(hass, entry)
     await guardian.async_update()
     hass.data[DOMAIN][DATA_CLIENT][entry.entry_id] = guardian
@@ -97,69 +68,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
             hass.config_entries.async_forward_entry_setup(entry, component)
         )
 
-    @_verify_domain_control
-    async def disable_ap(call):
-        """Disable the device's onboard access point."""
-        try:
-            async with guardian.client:
-                await guardian.client.wifi.disable_ap()
-        except GuardianError as err:
-            LOGGER.error("Error during service call: %s", err)
-            return
-
-    @_verify_domain_control
-    async def enable_ap(call):
-        """Enable the device's onboard access point."""
-        try:
-            async with guardian.client:
-                await guardian.client.wifi.enable_ap()
-        except GuardianError as err:
-            LOGGER.error("Error during service call: %s", err)
-            return
-
-    @_verify_domain_control
-    async def reboot(call):
-        """Reboot the device."""
-        try:
-            async with guardian.client:
-                await guardian.client.system.reboot()
-        except GuardianError as err:
-            LOGGER.error("Error during service call: %s", err)
-            return
-
-    @_verify_domain_control
-    async def reset_valve_diagnostics(call):
-        """Fully reset system motor diagnostics."""
-        try:
-            async with guardian.client:
-                await guardian.client.valve.reset()
-        except GuardianError as err:
-            LOGGER.error("Error during service call: %s", err)
-            return
-
-    @_verify_domain_control
-    async def upgrade_firmware(call):
-        """Upgrade the device firmware."""
-        try:
-            async with guardian.client:
-                await guardian.client.system.upgrade_firmware(
-                    url=call.data[CONF_URL],
-                    port=call.data[CONF_PORT],
-                    filename=call.data[CONF_FILENAME],
-                )
-        except GuardianError as err:
-            LOGGER.error("Error during service call: %s", err)
-            return
-
-    for service, method, schema in [
-        ("disable_ap", disable_ap, None),
-        ("enable_ap", enable_ap, None),
-        ("reboot", reboot, None),
-        ("reset_valve_diagnostics", reset_valve_diagnostics, None),
-        ("upgrade_firmware", upgrade_firmware, SERVICE_UPGRADE_FIRMWARE_SCHEMA),
-    ]:
-        async_register_admin_service(hass, DOMAIN, service, method, schema=schema)
-
     return True
 
 

homeassistant/components/guardian/switch.py

@@ -1,8 +1,16 @@
 """Switches for the Elexa Guardian integration."""
+from aioguardian.commands.system import (
+    DEFAULT_FIRMWARE_UPGRADE_FILENAME,
+    DEFAULT_FIRMWARE_UPGRADE_PORT,
+    DEFAULT_FIRMWARE_UPGRADE_URL,
+)
 from aioguardian.errors import GuardianError
+import voluptuous as vol
 
 from homeassistant.components.switch import SwitchEntity
+from homeassistant.const import CONF_FILENAME, CONF_PORT, CONF_URL
 from homeassistant.core import callback
+from homeassistant.helpers import config_validation as cv, entity_platform
 
 from . import Guardian, GuardianEntity
 from .const import DATA_CLIENT, DATA_VALVE_STATUS, DOMAIN, LOGGER, SWITCH_KIND_VALVE
@@ -12,10 +20,42 @@ ATTR_INST_CURRENT = "instantaneous_current"
 ATTR_INST_CURRENT_DDT = "instantaneous_current_ddt"
 ATTR_TRAVEL_COUNT = "travel_count"
 
+SERVICE_DISABLE_AP = "disable_ap"
+SERVICE_ENABLE_AP = "enable_ap"
+SERVICE_REBOOT = "reboot"
+SERVICE_RESET_VALVE_DIAGNOSTICS = "reset_valve_diagnostics"
+SERVICE_UPGRADE_FIRMWARE = "upgrade_firmware"
+
+SERVICE_UPGRADE_FIRMWARE_SCHEMA = vol.Schema(
+    {
+        vol.Optional(CONF_URL, default=DEFAULT_FIRMWARE_UPGRADE_URL): cv.url,
+        vol.Optional(CONF_PORT, default=DEFAULT_FIRMWARE_UPGRADE_PORT): cv.port,
+        vol.Optional(
+            CONF_FILENAME, default=DEFAULT_FIRMWARE_UPGRADE_FILENAME
+        ): cv.string,
+    }
+)
+
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up Guardian switches based on a config entry."""
     guardian = hass.data[DOMAIN][DATA_CLIENT][entry.entry_id]
+
+    platform = entity_platform.current_platform.get()
+
+    for service_name, schema, method in [
+        (SERVICE_DISABLE_AP, None, "async_disable_ap"),
+        (SERVICE_ENABLE_AP, None, "async_enable_ap"),
+        (SERVICE_REBOOT, None, "async_reboot"),
+        (SERVICE_RESET_VALVE_DIAGNOSTICS, None, "async_reset_valve_diagnostics"),
+        (
+            SERVICE_UPGRADE_FIRMWARE,
+            SERVICE_UPGRADE_FIRMWARE_SCHEMA,
+            "async_upgrade_firmware",
+        ),
+    ]:
+        platform.async_register_entity_service(service_name, schema, method)
+
     async_add_entities([GuardianSwitch(guardian)], True)
 
 
@@ -60,6 +100,48 @@ class GuardianSwitch(GuardianEntity, SwitchEntity):
             }
         )
 
+    async def async_disable_ap(self):
+        """Disable the device's onboard access point."""
+        try:
+            async with self._guardian.client:
+                await self._guardian.client.wifi.disable_ap()
+        except GuardianError as err:
+            LOGGER.error("Error during service call: %s", err)
+
+    async def async_enable_ap(self):
+        """Enable the device's onboard access point."""
+        try:
+            async with self._guardian.client:
+                await self._guardian.client.wifi.enable_ap()
+        except GuardianError as err:
+            LOGGER.error("Error during service call: %s", err)
+
+    async def async_reboot(self):
+        """Reboot the device."""
+        try:
+            async with self._guardian.client:
+                await self._guardian.client.system.reboot()
+        except GuardianError as err:
+            LOGGER.error("Error during service call: %s", err)
+
+    async def async_reset_valve_diagnostics(self):
+        """Fully reset system motor diagnostics."""
+        try:
+            async with self._guardian.client:
+                await self._guardian.client.valve.reset()
+        except GuardianError as err:
+            LOGGER.error("Error during service call: %s", err)
+
+    async def async_upgrade_firmware(self, *, url, port, filename):
+        """Upgrade the device firmware."""
+        try:
+            async with self._guardian.client:
+                await self._guardian.client.system.upgrade_firmware(
+                    url=url, port=port, filename=filename,
+                )
+        except GuardianError as err:
+            LOGGER.error("Error during service call: %s", err)
+
     async def async_turn_off(self, **kwargs) -> None:
         """Turn the valve off (closed)."""
         try: