Move neato service registration (#162146)

homeassistant/components/neato/__init__.py

@@ -10,18 +10,22 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_TOKEN, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.config_entry_oauth2_flow import (
     ImplementationUnavailableError,
     OAuth2Session,
     async_get_config_entry_implementation,
 )
+from homeassistant.helpers.typing import ConfigType
 
 from . import api
-from .const import NEATO_DOMAIN, NEATO_LOGIN
+from .const import DOMAIN, NEATO_LOGIN
 from .hub import NeatoHub
+from .services import async_setup_services
 
 _LOGGER = logging.getLogger(__name__)
 
+CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
 PLATFORMS = [
     Platform.BUTTON,
     Platform.CAMERA,
@@ -31,9 +35,15 @@ PLATFORMS = [
 ]
 
 
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up the component."""
+    async_setup_services(hass)
+    return True
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up config entry."""
-    hass.data.setdefault(NEATO_DOMAIN, {})
+    hass.data.setdefault(DOMAIN, {})
     if CONF_TOKEN not in entry.data:
         raise ConfigEntryAuthFailed
 
@@ -41,7 +51,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         implementation = await async_get_config_entry_implementation(hass, entry)
     except ImplementationUnavailableError as err:
         raise ConfigEntryNotReady(
-            translation_domain=NEATO_DOMAIN,
+            translation_domain=DOMAIN,
             translation_key="oauth2_implementation_unavailable",
         ) from err
 
@@ -55,7 +65,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise ConfigEntryNotReady from ex
 
     neato_session = api.ConfigEntryAuth(hass, entry, implementation)
-    hass.data[NEATO_DOMAIN][entry.entry_id] = neato_session
+    hass.data[DOMAIN][entry.entry_id] = neato_session
     hub = NeatoHub(hass, Account(neato_session))
 
     await hub.async_update_entry_unique_id(entry)
@@ -77,6 +87,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
-        hass.data[NEATO_DOMAIN].pop(entry.entry_id)
+        hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok

homeassistant/components/neato/config_flow.py

@@ -9,15 +9,15 @@ from typing import Any
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigFlowResult
 from homeassistant.helpers import config_entry_oauth2_flow
 
-from .const import NEATO_DOMAIN
+from .const import DOMAIN
 
 
 class OAuth2FlowHandler(
-    config_entry_oauth2_flow.AbstractOAuth2FlowHandler, domain=NEATO_DOMAIN
+    config_entry_oauth2_flow.AbstractOAuth2FlowHandler, domain=DOMAIN
 ):
     """Config flow to handle Neato Botvac OAuth2 authentication."""
 
-    DOMAIN = NEATO_DOMAIN
+    DOMAIN = DOMAIN
 
     @property
     def logger(self) -> logging.Logger:

homeassistant/components/neato/const.py

@@ -1,6 +1,6 @@
 """Constants for Neato integration."""
 
-NEATO_DOMAIN = "neato"
+DOMAIN = "neato"
 
 CONF_VENDOR = "vendor"
 NEATO_LOGIN = "neato_login"

homeassistant/components/neato/entity.py

@@ -7,7 +7,7 @@ from pybotvac import Robot
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity
 
-from .const import NEATO_DOMAIN
+from .const import DOMAIN
 
 
 class NeatoEntity(Entity):
@@ -19,6 +19,6 @@ class NeatoEntity(Entity):
         """Initialize Neato entity."""
         self.robot = robot
         self._attr_device_info: DeviceInfo = DeviceInfo(
-            identifiers={(NEATO_DOMAIN, self.robot.serial)},
+            identifiers={(DOMAIN, self.robot.serial)},
             name=self.robot.name,
         )

homeassistant/components/neato/services.py

@@ -0,0 +1,36 @@
+"""Neato services."""
+
+from __future__ import annotations
+
+import voluptuous as vol
+
+from homeassistant.components.vacuum import DOMAIN as VACUUM_DOMAIN
+from homeassistant.const import ATTR_MODE
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import config_validation as cv, service
+
+from .const import DOMAIN
+
+ATTR_NAVIGATION = "navigation"
+ATTR_CATEGORY = "category"
+ATTR_ZONE = "zone"
+
+
+@callback
+def async_setup_services(hass: HomeAssistant) -> None:
+    """Set up services."""
+
+    # Vacuum Services
+    service.async_register_platform_entity_service(
+        hass,
+        DOMAIN,
+        "custom_cleaning",
+        entity_domain=VACUUM_DOMAIN,
+        schema={
+            vol.Optional(ATTR_MODE, default=2): cv.positive_int,
+            vol.Optional(ATTR_NAVIGATION, default=1): cv.positive_int,
+            vol.Optional(ATTR_CATEGORY, default=4): cv.positive_int,
+            vol.Optional(ATTR_ZONE): cv.string,
+        },
+        func="neato_custom_cleaning",
+    )

homeassistant/components/neato/vacuum.py

@@ -8,7 +8,6 @@ from typing import Any
 
 from pybotvac import Robot
 from pybotvac.exceptions import NeatoRobotException
-import voluptuous as vol
 
 from homeassistant.components.vacuum import (
     ATTR_STATUS,
@@ -17,9 +16,7 @@ from homeassistant.components.vacuum import (
     VacuumEntityFeature,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import ATTR_MODE
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import config_validation as cv, entity_platform
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
@@ -52,10 +49,6 @@ ATTR_CLEAN_PAUSE_TIME = "clean_pause_time"
 ATTR_CLEAN_ERROR_TIME = "clean_error_time"
 ATTR_LAUNCHED_FROM = "launched_from"
 
-ATTR_NAVIGATION = "navigation"
-ATTR_CATEGORY = "category"
-ATTR_ZONE = "zone"
-
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -77,20 +70,6 @@ async def async_setup_entry(
     _LOGGER.debug("Adding vacuums %s", dev)
     async_add_entities(dev, True)
 
-    platform = entity_platform.async_get_current_platform()
-    assert platform is not None
-
-    platform.async_register_entity_service(
-        "custom_cleaning",
-        {
-            vol.Optional(ATTR_MODE, default=2): cv.positive_int,
-            vol.Optional(ATTR_NAVIGATION, default=1): cv.positive_int,
-            vol.Optional(ATTR_CATEGORY, default=4): cv.positive_int,
-            vol.Optional(ATTR_ZONE): cv.string,
-        },
-        "neato_custom_cleaning",
-    )
-
 
 class NeatoConnectedVacuum(NeatoEntity, StateVacuumEntity):
     """Representation of a Neato Connected Vacuum."""

tests/components/neato/test_config_flow.py

@@ -10,7 +10,7 @@ from homeassistant.components.application_credentials import (
     ClientCredential,
     async_import_client_credential,
 )
-from homeassistant.components.neato.const import NEATO_DOMAIN
+from homeassistant.components.neato.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
@@ -37,7 +37,7 @@ async def test_full_flow(
     """Check full flow."""
     assert await setup.async_setup_component(hass, "neato", {})
     await async_import_client_credential(
-        hass, NEATO_DOMAIN, ClientCredential(CLIENT_ID, CLIENT_SECRET)
+        hass, DOMAIN, ClientCredential(CLIENT_ID, CLIENT_SECRET)
     )
 
     result = await hass.config_entries.flow.async_init(
@@ -79,14 +79,14 @@ async def test_full_flow(
     ) as mock_setup:
         await hass.config_entries.flow.async_configure(result["flow_id"])
 
-    assert len(hass.config_entries.async_entries(NEATO_DOMAIN)) == 1
+    assert len(hass.config_entries.async_entries(DOMAIN)) == 1
     assert len(mock_setup.mock_calls) == 1
 
 
 async def test_abort_if_already_setup(hass: HomeAssistant) -> None:
     """Test we abort if Neato is already setup."""
     entry = MockConfigEntry(
-        domain=NEATO_DOMAIN,
+        domain=DOMAIN,
         data={"auth_implementation": "neato", "token": {"some": "data"}},
     )
     entry.add_to_hass(hass)
@@ -108,12 +108,12 @@ async def test_reauth(
     """Test initialization of the reauth flow."""
     assert await setup.async_setup_component(hass, "neato", {})
     await async_import_client_credential(
-        hass, NEATO_DOMAIN, ClientCredential(CLIENT_ID, CLIENT_SECRET)
+        hass, DOMAIN, ClientCredential(CLIENT_ID, CLIENT_SECRET)
     )
 
     entry = MockConfigEntry(
         entry_id="my_entry",
-        domain=NEATO_DOMAIN,
+        domain=DOMAIN,
         data={"username": "abcdef", "password": "123456", "vendor": "neato"},
     )
     entry.add_to_hass(hass)
@@ -160,5 +160,5 @@ async def test_reauth(
     assert result3["type"] is FlowResultType.ABORT
     assert result3["reason"] == "reauth_successful"
     assert new_entry.state is ConfigEntryState.LOADED
-    assert len(hass.config_entries.async_entries(NEATO_DOMAIN)) == 1
+    assert len(hass.config_entries.async_entries(DOMAIN)) == 1
     assert len(mock_setup.mock_calls) == 1

tests/components/neato/test_init.py

@@ -2,7 +2,7 @@
 
 from unittest.mock import patch
 
-from homeassistant.components.neato.const import NEATO_DOMAIN
+from homeassistant.components.neato.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.config_entry_oauth2_flow import (
@@ -17,9 +17,9 @@ async def test_oauth_implementation_not_available(
 ) -> None:
     """Test that unavailable OAuth implementation raises ConfigEntryNotReady."""
     config_entry = MockConfigEntry(
-        domain=NEATO_DOMAIN,
+        domain=DOMAIN,
         data={
-            "auth_implementation": NEATO_DOMAIN,
+            "auth_implementation": DOMAIN,
             "token": {
                 "refresh_token": "mock-refresh-token",
                 "access_token": "mock-access-token",