From 3ba1fbe69d24f7d86ca4a8127d23bccf0c8e4427 Mon Sep 17 00:00:00 2001
From: Nathan Spencer <natekspencer@gmail.com>
Date: Tue, 30 Aug 2022 20:01:49 -0600
Subject: [PATCH] Add pet weight sensor for Litter-Robot 4 (#77566)

---
 .../components/litterrobot/sensor.py          | 115 +++++++++---------
 1 file changed, 59 insertions(+), 56 deletions(-)

diff --git a/homeassistant/components/litterrobot/sensor.py b/homeassistant/components/litterrobot/sensor.py
index 386d0e04f3c..90bdfcbda73 100644
--- a/homeassistant/components/litterrobot/sensor.py
+++ b/homeassistant/components/litterrobot/sensor.py
@@ -1,13 +1,12 @@
 """Support for Litter-Robot sensors."""
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable
+from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime
-import itertools
 from typing import Any, Generic, Union, cast
 
-from pylitterbot import FeederRobot, LitterRobot
+from pylitterbot import FeederRobot, LitterRobot, LitterRobot4, Robot
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -15,7 +14,7 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import PERCENTAGE
+from homeassistant.const import MASS_POUNDS, PERCENTAGE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -77,45 +76,56 @@ class LitterRobotSensorEntity(LitterRobotEntity[_RobotT], SensorEntity):
         return super().icon
 
 
-LITTER_ROBOT_SENSORS = [
-    RobotSensorEntityDescription[LitterRobot](
-        name="Waste Drawer",
-        key="waste_drawer_level",
-        native_unit_of_measurement=PERCENTAGE,
-        icon_fn=lambda state: icon_for_gauge_level(state, 10),
-    ),
-    RobotSensorEntityDescription[LitterRobot](
-        name="Sleep Mode Start Time",
-        key="sleep_mode_start_time",
-        device_class=SensorDeviceClass.TIMESTAMP,
-        should_report=lambda robot: robot.sleep_mode_enabled,
-    ),
-    RobotSensorEntityDescription[LitterRobot](
-        name="Sleep Mode End Time",
-        key="sleep_mode_end_time",
-        device_class=SensorDeviceClass.TIMESTAMP,
-        should_report=lambda robot: robot.sleep_mode_enabled,
-    ),
-    RobotSensorEntityDescription[LitterRobot](
-        name="Last Seen",
-        key="last_seen",
-        device_class=SensorDeviceClass.TIMESTAMP,
-        entity_category=EntityCategory.DIAGNOSTIC,
-    ),
-    RobotSensorEntityDescription[LitterRobot](
-        name="Status Code",
-        key="status_code",
-        device_class="litterrobot__status_code",
-        entity_category=EntityCategory.DIAGNOSTIC,
-    ),
-]
-
-FEEDER_ROBOT_SENSOR = RobotSensorEntityDescription[FeederRobot](
-    name="Food Level",
-    key="food_level",
-    native_unit_of_measurement=PERCENTAGE,
-    icon_fn=lambda state: icon_for_gauge_level(state, 10),
-)
+ROBOT_SENSOR_MAP: dict[type[Robot], list[RobotSensorEntityDescription]] = {
+    LitterRobot: [
+        RobotSensorEntityDescription[LitterRobot](
+            name="Waste Drawer",
+            key="waste_drawer_level",
+            native_unit_of_measurement=PERCENTAGE,
+            icon_fn=lambda state: icon_for_gauge_level(state, 10),
+        ),
+        RobotSensorEntityDescription[LitterRobot](
+            name="Sleep Mode Start Time",
+            key="sleep_mode_start_time",
+            device_class=SensorDeviceClass.TIMESTAMP,
+            should_report=lambda robot: robot.sleep_mode_enabled,
+        ),
+        RobotSensorEntityDescription[LitterRobot](
+            name="Sleep Mode End Time",
+            key="sleep_mode_end_time",
+            device_class=SensorDeviceClass.TIMESTAMP,
+            should_report=lambda robot: robot.sleep_mode_enabled,
+        ),
+        RobotSensorEntityDescription[LitterRobot](
+            name="Last Seen",
+            key="last_seen",
+            device_class=SensorDeviceClass.TIMESTAMP,
+            entity_category=EntityCategory.DIAGNOSTIC,
+        ),
+        RobotSensorEntityDescription[LitterRobot](
+            name="Status Code",
+            key="status_code",
+            device_class="litterrobot__status_code",
+            entity_category=EntityCategory.DIAGNOSTIC,
+        ),
+    ],
+    LitterRobot4: [
+        RobotSensorEntityDescription[LitterRobot4](
+            key="pet_weight",
+            name="Pet weight",
+            icon="mdi:scale",
+            native_unit_of_measurement=MASS_POUNDS,
+        )
+    ],
+    FeederRobot: [
+        RobotSensorEntityDescription[FeederRobot](
+            name="Food level",
+            key="food_level",
+            native_unit_of_measurement=PERCENTAGE,
+            icon_fn=lambda state: icon_for_gauge_level(state, 10),
+        )
+    ],
+}
 
 
 async def async_setup_entry(
@@ -125,17 +135,10 @@ async def async_setup_entry(
 ) -> None:
     """Set up Litter-Robot sensors using config entry."""
     hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
-    entities: Iterable[LitterRobotSensorEntity] = itertools.chain(
-        (
-            LitterRobotSensorEntity(robot=robot, hub=hub, description=description)
-            for description in LITTER_ROBOT_SENSORS
-            for robot in hub.litter_robots()
-        ),
-        (
-            LitterRobotSensorEntity(
-                robot=robot, hub=hub, description=FEEDER_ROBOT_SENSOR
-            )
-            for robot in hub.feeder_robots()
-        ),
+    async_add_entities(
+        LitterRobotSensorEntity(robot=robot, hub=hub, description=description)
+        for robot in hub.account.robots
+        for robot_type, entity_descriptions in ROBOT_SENSOR_MAP.items()
+        if isinstance(robot, robot_type)
+        for description in entity_descriptions
     )
-    async_add_entities(entities)