From 538de6d1f316bffee2b1bd40cea528ce61aa1f14 Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Mon, 21 Aug 2023 12:04:12 +0200 Subject: [PATCH] Introduce base class for Neato (#98071) Co-authored-by: Robert Resch --- .coveragerc | 1 + homeassistant/components/neato/button.py | 13 ++++-------- homeassistant/components/neato/camera.py | 25 +++++----------------- homeassistant/components/neato/entity.py | 27 ++++++++++++++++++++++++ homeassistant/components/neato/sensor.py | 18 ++++------------ homeassistant/components/neato/switch.py | 17 ++++----------- homeassistant/components/neato/vacuum.py | 21 ++++++++---------- 7 files changed, 54 insertions(+), 68 deletions(-) create mode 100644 homeassistant/components/neato/entity.py diff --git a/.coveragerc b/.coveragerc index 58aa8eb52c3..2fbfa15997b 100644 --- a/.coveragerc +++ b/.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 diff --git a/homeassistant/components/neato/button.py b/homeassistant/components/neato/button.py index 4bbd9196932..8b23bbe4681 100644 --- a/homeassistant/components/neato/button.py +++ b/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.""" diff --git a/homeassistant/components/neato/camera.py b/homeassistant/components/neato/camera.py index 5b13d12d37f..c1513bb1de6 100644 --- a/homeassistant/components/neato/camera.py +++ b/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 ( - NEATO_DOMAIN, - NEATO_LOGIN, - NEATO_MAP_DATA, - NEATO_ROBOTS, - SCAN_INTERVAL_MINUTES, -) +from .const import NEATO_LOGIN, NEATO_MAP_DATA, NEATO_ROBOTS, SCAN_INTERVAL_MINUTES +from .entity import NeatoEntity 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__() - self.robot = robot + super().__init__(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.""" diff --git a/homeassistant/components/neato/entity.py b/homeassistant/components/neato/entity.py new file mode 100644 index 00000000000..43072f19693 --- /dev/null +++ b/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, + ) diff --git a/homeassistant/components/neato/sensor.py b/homeassistant/components/neato/sensor.py index 60aeb52af05..452f1bc3a9c 100644 --- a/homeassistant/components/neato/sensor.py +++ b/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, - ) diff --git a/homeassistant/components/neato/switch.py b/homeassistant/components/neato/switch.py index f6d159fcc1b..a80d05eef23 100644 --- a/homeassistant/components/neato/switch.py +++ b/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: diff --git a/homeassistant/components/neato/vacuum.py b/homeassistant/components/neato/vacuum.py index f70e79f3fc0..ecc39e515c2 100644 --- a/homeassistant/components/neato/vacuum.py +++ b/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 - return DeviceInfo( - identifiers={(NEATO_DOMAIN, self._robot_serial)}, - manufacturer=stats["battery"]["vendor"] if stats else None, - model=stats["model"] if stats else None, - name=self.robot.name, - sw_version=stats["firmware"] if stats else None, - ) + device_info = super().device_info + if self._robot_stats: + device_info["manufacturer"] = self._robot_stats["battery"]["vendor"] + device_info["model"] = self._robot_stats["model"] + device_info["sw_version"] = self._robot_stats["firmware"] + return device_info def start(self) -> None: """Start cleaning or resume cleaning."""