Introduce base class for Neato (#98071)

Co-authored-by: Robert Resch <robert@resch.dev>

.coveragerc

@@ -781,6 +781,7 @@ omit =
     homeassistant/components/neato/__init__.py
     homeassistant/components/neato/api.py
     homeassistant/components/neato/camera.py
+    homeassistant/components/neato/entity.py
     homeassistant/components/neato/hub.py
     homeassistant/components/neato/sensor.py
     homeassistant/components/neato/switch.py

homeassistant/components/neato/button.py

@@ -7,10 +7,10 @@ from homeassistant.components.button import ButtonEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import NEATO_DOMAIN, NEATO_ROBOTS
+from .const import NEATO_ROBOTS
+from .entity import NeatoEntity
 
 
 async def async_setup_entry(
@@ -22,10 +22,9 @@ async def async_setup_entry(
     async_add_entities(entities, True)
 
 
-class NeatoDismissAlertButton(ButtonEntity):
+class NeatoDismissAlertButton(NeatoEntity, ButtonEntity):
     """Representation of a dismiss_alert button entity."""
 
-    _attr_has_entity_name = True
     _attr_translation_key = "dismiss_alert"
     _attr_entity_category = EntityCategory.CONFIG
 
@@ -34,12 +33,8 @@ class NeatoDismissAlertButton(ButtonEntity):
         robot: Robot,
     ) -> None:
         """Initialize a dismiss_alert Neato button entity."""
-        self.robot = robot
+        super().__init__(robot)
         self._attr_unique_id = f"{robot.serial}_dismiss_alert"
-        self._attr_device_info = DeviceInfo(
-            identifiers={(NEATO_DOMAIN, robot.serial)},
-            name=robot.name,
-        )
 
     async def async_press(self) -> None:
         """Press the button."""

homeassistant/components/neato/camera.py

@@ -12,16 +12,10 @@ from urllib3.response import HTTPResponse
 from homeassistant.components.camera import Camera
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import (
+from .const import NEATO_LOGIN, NEATO_MAP_DATA, NEATO_ROBOTS, SCAN_INTERVAL_MINUTES
-    NEATO_DOMAIN,
+from .entity import NeatoEntity
-    NEATO_LOGIN,
-    NEATO_MAP_DATA,
-    NEATO_ROBOTS,
-    SCAN_INTERVAL_MINUTES,
-)
 from .hub import NeatoHub
 
 _LOGGER = logging.getLogger(__name__)
@@ -48,18 +42,17 @@ async def async_setup_entry(
     async_add_entities(dev, True)
 
 
-class NeatoCleaningMap(Camera):
+class NeatoCleaningMap(NeatoEntity, Camera):
     """Neato cleaning map for last clean."""
 
-    _attr_has_entity_name = True
     _attr_translation_key = "cleaning_map"
 
     def __init__(
         self, neato: NeatoHub, robot: Robot, mapdata: dict[str, Any] | None
     ) -> None:
         """Initialize Neato cleaning map."""
-        super().__init__()
+        super().__init__(robot)
-        self.robot = robot
+        Camera.__init__(self)
         self.neato = neato
         self._mapdata = mapdata
         self._available = neato is not None
@@ -126,14 +119,6 @@ class NeatoCleaningMap(Camera):
         """Return if the robot is available."""
         return self._available
 
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Device info for neato robot."""
-        return DeviceInfo(
-            identifiers={(NEATO_DOMAIN, self._robot_serial)},
-            name=self.robot.name,
-        )
-
     @property
     def extra_state_attributes(self) -> dict[str, Any]:
         """Return the state attributes of the vacuum cleaner."""

homeassistant/components/neato/entity.py

@@ -0,0 +1,27 @@
+"""Base entity for Neato."""
+from __future__ import annotations
+
+from pybotvac import Robot
+
+from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers.entity import Entity
+
+from .const import NEATO_DOMAIN
+
+
+class NeatoEntity(Entity):
+    """Base Neato entity."""
+
+    _attr_has_entity_name = True
+
+    def __init__(self, robot: Robot) -> None:
+        """Initialize Neato entity."""
+        self.robot = robot
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device info."""
+        return DeviceInfo(
+            identifiers={(NEATO_DOMAIN, self.robot.serial)},
+            name=self.robot.name,
+        )

homeassistant/components/neato/sensor.py

@@ -12,10 +12,10 @@ from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import NEATO_DOMAIN, NEATO_LOGIN, NEATO_ROBOTS, SCAN_INTERVAL_MINUTES
+from .const import NEATO_LOGIN, NEATO_ROBOTS, SCAN_INTERVAL_MINUTES
+from .entity import NeatoEntity
 from .hub import NeatoHub
 
 _LOGGER = logging.getLogger(__name__)
@@ -41,14 +41,12 @@ async def async_setup_entry(
     async_add_entities(dev, True)
 
 
-class NeatoSensor(SensorEntity):
+class NeatoSensor(NeatoEntity, SensorEntity):
     """Neato sensor."""
 
-    _attr_has_entity_name = True
-
     def __init__(self, neato: NeatoHub, robot: Robot) -> None:
         """Initialize Neato sensor."""
-        self.robot = robot
+        super().__init__(robot)
         self._available: bool = False
         self._robot_serial: str = self.robot.serial
         self._state: dict[str, Any] | None = None
@@ -100,11 +98,3 @@ class NeatoSensor(SensorEntity):
     def native_unit_of_measurement(self) -> str:
         """Return unit of measurement."""
         return PERCENTAGE
-
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Device info for neato robot."""
-        return DeviceInfo(
-            identifiers={(NEATO_DOMAIN, self._robot_serial)},
-            name=self.robot.name,
-        )

homeassistant/components/neato/switch.py

@@ -12,10 +12,10 @@ from homeassistant.components.switch import SwitchEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import STATE_OFF, STATE_ON, EntityCategory
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import NEATO_DOMAIN, NEATO_LOGIN, NEATO_ROBOTS, SCAN_INTERVAL_MINUTES
+from .const import NEATO_LOGIN, NEATO_ROBOTS, SCAN_INTERVAL_MINUTES
+from .entity import NeatoEntity
 from .hub import NeatoHub
 
 _LOGGER = logging.getLogger(__name__)
@@ -45,16 +45,15 @@ async def async_setup_entry(
     async_add_entities(dev, True)
 
 
-class NeatoConnectedSwitch(SwitchEntity):
+class NeatoConnectedSwitch(NeatoEntity, SwitchEntity):
     """Neato Connected Switches."""
 
-    _attr_has_entity_name = True
     _attr_translation_key = "schedule"
 
     def __init__(self, neato: NeatoHub, robot: Robot, switch_type: str) -> None:
         """Initialize the Neato Connected switches."""
+        super().__init__(robot)
         self.type = switch_type
-        self.robot = robot
         self._available = False
         self._state: dict[str, Any] | None = None
         self._schedule_state: str | None = None
@@ -109,14 +108,6 @@ class NeatoConnectedSwitch(SwitchEntity):
         """Device entity category."""
         return EntityCategory.CONFIG
 
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Device info for neato robot."""
-        return DeviceInfo(
-            identifiers={(NEATO_DOMAIN, self._robot_serial)},
-            name=self.robot.name,
-        )
-
     def turn_on(self, **kwargs: Any) -> None:
         """Turn the switch on."""
         if self.type == SWITCH_TYPE_SCHEDULE:

homeassistant/components/neato/vacuum.py

@@ -30,13 +30,13 @@ from .const import (
     ALERTS,
     ERRORS,
     MODE,
-    NEATO_DOMAIN,
     NEATO_LOGIN,
     NEATO_MAP_DATA,
     NEATO_PERSISTENT_MAPS,
     NEATO_ROBOTS,
     SCAN_INTERVAL_MINUTES,
 )
+from .entity import NeatoEntity
 from .hub import NeatoHub
 
 _LOGGER = logging.getLogger(__name__)
@@ -91,7 +91,7 @@ async def async_setup_entry(
     )
 
 
-class NeatoConnectedVacuum(StateVacuumEntity):
+class NeatoConnectedVacuum(NeatoEntity, StateVacuumEntity):
     """Representation of a Neato Connected Vacuum."""
 
     _attr_icon = "mdi:robot-vacuum-variant"
@@ -106,7 +106,6 @@ class NeatoConnectedVacuum(StateVacuumEntity):
         | VacuumEntityFeature.MAP
         | VacuumEntityFeature.LOCATE
     )
-    _attr_has_entity_name = True
     _attr_name = None
 
     def __init__(
@@ -117,7 +116,7 @@ class NeatoConnectedVacuum(StateVacuumEntity):
         persistent_maps: dict[str, Any] | None,
     ) -> None:
         """Initialize the Neato Connected Vacuum."""
-        self.robot = robot
+        super().__init__(robot)
         self._attr_available: bool = neato is not None
         self._mapdata = mapdata
         self._robot_has_map: bool = self.robot.has_persistent_maps
@@ -300,14 +299,12 @@ class NeatoConnectedVacuum(StateVacuumEntity):
     @property
     def device_info(self) -> DeviceInfo:
         """Device info for neato robot."""
-        stats = self._robot_stats
+        device_info = super().device_info
-        return DeviceInfo(
+        if self._robot_stats:
-            identifiers={(NEATO_DOMAIN, self._robot_serial)},
+            device_info["manufacturer"] = self._robot_stats["battery"]["vendor"]
-            manufacturer=stats["battery"]["vendor"] if stats else None,
+            device_info["model"] = self._robot_stats["model"]
-            model=stats["model"] if stats else None,
+            device_info["sw_version"] = self._robot_stats["firmware"]
-            name=self.robot.name,
+        return device_info
-            sw_version=stats["firmware"] if stats else None,
-        )
 
     def start(self) -> None:
         """Start cleaning or resume cleaning."""