From fe881230dbd7631441864ebb1aef2f19b7c832cd Mon Sep 17 00:00:00 2001
From: Nathan Spencer <natekspencer@gmail.com>
Date: Tue, 30 Aug 2022 12:30:42 -0600
Subject: [PATCH] Add support for Feeder-Robot button (#77501)

Co-authored-by: epenet <6771947+epenet@users.noreply.github.com>
---
 .../components/litterrobot/button.py          | 81 +++++++++++++++----
 1 file changed, 65 insertions(+), 16 deletions(-)

diff --git a/homeassistant/components/litterrobot/button.py b/homeassistant/components/litterrobot/button.py
index b833500ec4c..74c659fd474 100644
--- a/homeassistant/components/litterrobot/button.py
+++ b/homeassistant/components/litterrobot/button.py
@@ -1,20 +1,23 @@
 """Support for Litter-Robot button."""
 from __future__ import annotations
 
-from pylitterbot import LitterRobot3
+from collections.abc import Callable, Coroutine, Iterable
+from dataclasses import dataclass
+import itertools
+from typing import Any, Generic
 
-from homeassistant.components.button import ButtonEntity
+from pylitterbot import FeederRobot, LitterRobot3
+
+from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .entity import LitterRobotEntity
+from .entity import LitterRobotEntity, _RobotT
 from .hub import LitterRobotHub
 
-TYPE_RESET_WASTE_DRAWER = "Reset Waste Drawer"
-
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -23,22 +26,68 @@ async def async_setup_entry(
 ) -> None:
     """Set up Litter-Robot cleaner using config entry."""
     hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities(
-        LitterRobotResetWasteDrawerButton(
-            robot=robot, entity_type=TYPE_RESET_WASTE_DRAWER, hub=hub
-        )
-        for robot in hub.litter_robots()
-        if isinstance(robot, LitterRobot3)
+    entities: Iterable[LitterRobotButtonEntity] = itertools.chain(
+        (
+            LitterRobotButtonEntity(
+                robot=robot, hub=hub, description=LITTER_ROBOT_BUTTON
+            )
+            for robot in hub.litter_robots()
+            if isinstance(robot, LitterRobot3)
+        ),
+        (
+            LitterRobotButtonEntity(
+                robot=robot, hub=hub, description=FEEDER_ROBOT_BUTTON
+            )
+            for robot in hub.feeder_robots()
+        ),
     )
+    async_add_entities(entities)
 
 
-class LitterRobotResetWasteDrawerButton(LitterRobotEntity[LitterRobot3], ButtonEntity):
-    """Litter-Robot reset waste drawer button."""
+@dataclass
+class RequiredKeysMixin(Generic[_RobotT]):
+    """A class that describes robot button entity required keys."""
 
-    _attr_icon = "mdi:delete-variant"
-    _attr_entity_category = EntityCategory.CONFIG
+    press_fn: Callable[[_RobotT], Coroutine[Any, Any, bool]]
+
+
+@dataclass
+class RobotButtonEntityDescription(ButtonEntityDescription, RequiredKeysMixin[_RobotT]):
+    """A class that describes robot button entities."""
+
+
+LITTER_ROBOT_BUTTON = RobotButtonEntityDescription[LitterRobot3](
+    key="reset_waste_drawer",
+    name="Reset Waste Drawer",
+    icon="mdi:delete-variant",
+    entity_category=EntityCategory.CONFIG,
+    press_fn=lambda robot: robot.reset_waste_drawer(),
+)
+FEEDER_ROBOT_BUTTON = RobotButtonEntityDescription[FeederRobot](
+    key="give_snack",
+    name="Give snack",
+    icon="mdi:candy-outline",
+    press_fn=lambda robot: robot.give_snack(),
+)
+
+
+class LitterRobotButtonEntity(LitterRobotEntity[_RobotT], ButtonEntity):
+    """Litter-Robot button entity."""
+
+    entity_description: RobotButtonEntityDescription[_RobotT]
+
+    def __init__(
+        self,
+        robot: _RobotT,
+        hub: LitterRobotHub,
+        description: RobotButtonEntityDescription[_RobotT],
+    ) -> None:
+        """Initialize a Litter-Robot button entity."""
+        assert description.name
+        super().__init__(robot, description.name, hub)
+        self.entity_description = description
 
     async def async_press(self) -> None:
         """Press the button."""
-        await self.robot.reset_waste_drawer()
+        await self.entity_description.press_fn(self.robot)
         self.coordinator.async_set_updated_data(True)