From cf736bae2a228b3c9f84d4031e0da1f5a73ea120 Mon Sep 17 00:00:00 2001
From: Nathan Spencer <natekspencer@gmail.com>
Date: Thu, 4 Apr 2024 23:59:03 -0600
Subject: [PATCH] Handle cat detected status from displayCode field

---
 pylitterbot/robot/litterrobot4.py | 24 +++++++++++++++++-------
 tests/common.py                   |  1 +
 tests/test_litterrobot4.py        |  2 ++
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/pylitterbot/robot/litterrobot4.py b/pylitterbot/robot/litterrobot4.py
index c8fb78a..00a8e84 100644
--- a/pylitterbot/robot/litterrobot4.py
+++ b/pylitterbot/robot/litterrobot4.py
@@ -54,6 +54,7 @@
     "CYCLE_STATE_CAT_DETECT": LitterBoxStatus.CAT_SENSOR_INTERRUPTED,
     "CYCLE_STATE_PAUSE": LitterBoxStatus.PAUSED,
 }
+DISPLAY_CODE_STATUS_MAP = {"DC_CAT_DETECT": LitterBoxStatus.CAT_DETECTED}
 
 LITTER_LEVEL_EMPTY = 500
 
@@ -246,16 +247,25 @@ def sleep_mode_end_time(self) -> datetime | None:
 
     @property
     def status(self) -> LitterBoxStatus:
-        """Return the status of the Litter-Robot."""
+        """Return the status of the Litter-Robot.
+
+        The Litter-Robot 4's status is determined based on the values of:
+          - `displayCode`
+          - `robotCycleState`
+          - `robotStatus`
+          - `isDFIFull`
+          - `isOnline`
+        """
         if not self.is_online:
             return LitterBoxStatus.OFFLINE
-        cycle_state = self._data["robotCycleState"]
-        status = self._data["robotStatus"]
-        if status == "ROBOT_IDLE" and self.is_waste_drawer_full:
+        if status := CYCLE_STATE_STATUS_MAP.get(self._data["robotCycleState"]):
+            return status
+        if status := DISPLAY_CODE_STATUS_MAP.get(self._data["displayCode"]):
+            return status
+        status = LR4_STATUS_MAP.get(self._data["robotStatus"], LitterBoxStatus.UNKNOWN)
+        if status == LitterBoxStatus.READY and self.is_waste_drawer_full:
             return LitterBoxStatus.DRAWER_FULL
-        return CYCLE_STATE_STATUS_MAP.get(
-            cycle_state, LR4_STATUS_MAP.get(status, LitterBoxStatus.UNKNOWN)
-        )
+        return status
 
     @property
     def status_code(self) -> str | None:
diff --git a/tests/common.py b/tests/common.py
index 2212b3d..5fa4729 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -168,6 +168,7 @@
     "firmwareUpdateStatus": "SUCCEEDED",
     "unitPowerType": "AC",
     "catWeight": 7.93,
+    "displayCode": "DC_MODE_IDLE",
     "unitTimezone": "America/Denver",
     "cleanCycleWaitTime": 7,
     "isKeypadLockout": False,
diff --git a/tests/test_litterrobot4.py b/tests/test_litterrobot4.py
index 454efea..d799d94 100644
--- a/tests/test_litterrobot4.py
+++ b/tests/test_litterrobot4.py
@@ -379,6 +379,7 @@ async def test_litter_robot_4_cleaning(mock_account: Account) -> None:
 @pytest.mark.parametrize(
     "updated_data,status",
     [
+        ({"isDFIFull": True}, LitterBoxStatus.DRAWER_FULL),
         ({"isOnline": False}, LitterBoxStatus.OFFLINE),
         (
             {"isOnline": False, "robotCycleState": "CYCLE_STATE_PAUSE"},
@@ -403,6 +404,7 @@ async def test_litter_robot_4_cleaning(mock_account: Account) -> None:
         ({"robotStatus": "ROBOT_POWER_DOWN"}, LitterBoxStatus.POWER_DOWN),
         ({"robotStatus": "ROBOT_POWER_OFF"}, LitterBoxStatus.OFF),
         ({"robotStatus": "ROBOT_POWER_UP"}, LitterBoxStatus.POWER_UP),
+        ({"displayCode": "DC_CAT_DETECT"}, LitterBoxStatus.CAT_DETECTED),
     ],
 )
 async def test_litter_robot_4_status(