From fd973a00c022170415001439824919ad450ae75e Mon Sep 17 00:00:00 2001
From: SLaks <Dev@SLaks.net>
Date: Tue, 19 Mar 2024 15:23:40 -0400
Subject: [PATCH] Roborock: Add WiFi strength sensor

I plan to use this to draw a WiFi signal strength map of my house
---
 miio/integrations/roborock/vacuum/vacuum.py    |  7 +++++++
 .../roborock/vacuum/vacuumcontainers.py        | 18 ++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/miio/integrations/roborock/vacuum/vacuum.py b/miio/integrations/roborock/vacuum/vacuum.py
index 2b294393f..8cd15be5f 100644
--- a/miio/integrations/roborock/vacuum/vacuum.py
+++ b/miio/integrations/roborock/vacuum/vacuum.py
@@ -47,6 +47,7 @@
     CarpetModeStatus,
     CleaningDetails,
     CleaningSummary,
+    ConnectivityStatus,
     ConsumableStatus,
     DNDStatus,
     MapList,
@@ -155,6 +156,7 @@ def __init__(
         self._map_enum_cache = None
         self._status_helper = UpdateHelper(self.vacuum_status)
         self._status_helper.add_update_method("consumables", self.consumable_status)
+        self._status_helper.add_update_method("connectivity", self.connectivity_status)
         self._status_helper.add_update_method("dnd_status", self.dnd_status)
         self._status_helper.add_update_method("clean_history", self.clean_history)
         self._status_helper.add_update_method("last_clean", self.last_clean_details)
@@ -620,6 +622,11 @@ def dnd_status(self) -> DNDStatus:
         #  'start_hour': 22, 'end_hour': 8}], 'id': 1}
         return DNDStatus(self.send("get_dnd_timer")[0])
 
+    @command()
+    def connectivity_status(self) -> ConnectivityStatus:
+        """Returns WiFi connectivity information."""
+        return ConnectivityStatus(self.info(skip_cache=True))
+
     @command(
         click.argument("start_hr", type=int),
         click.argument("start_min", type=int),
diff --git a/miio/integrations/roborock/vacuum/vacuumcontainers.py b/miio/integrations/roborock/vacuum/vacuumcontainers.py
index ea42451f2..e6d64f3a2 100644
--- a/miio/integrations/roborock/vacuum/vacuumcontainers.py
+++ b/miio/integrations/roborock/vacuum/vacuumcontainers.py
@@ -8,6 +8,7 @@
 from pytz import BaseTzInfo
 
 from miio.device import DeviceStatus
+from miio.deviceinfo import DeviceInfo
 from miio.devicestatus import sensor, setting
 from miio.identifiers import VacuumId, VacuumState
 from miio.utils import pretty_seconds, pretty_time
@@ -1062,3 +1063,20 @@ def enabled(self) -> bool:
     def dry_time(self) -> timedelta:
         """Return mop dry time."""
         return pretty_seconds(self.data["on"]["dry_time"])
+
+
+class ConnectivityStatus(DeviceStatus):
+    def __init__(self, info: DeviceInfo) -> None:
+        self.info = info
+
+    @property
+    @sensor(
+        "Wifi Signal Strengh",
+        unit="dBm",
+        icon="mdi:wifi-strength-2",  # TODO: Make this dynamic?
+        device_class="signal_strength",
+        entity_category="diagnostic",
+    )
+    def total_duration(self) -> timedelta:
+        """Total cleaning duration."""
+        return self.info.accesspoint["rssi"]