From f14df412766dfc4d444034efb7b4f346c21e6fcb Mon Sep 17 00:00:00 2001
From: Teemu Rytilahti <tpr@iki.fi>
Date: Sun, 15 Mar 2020 22:04:12 +0100
Subject: [PATCH] Add fan_speed_presets() for querying available fan speeds

This returns a dictionary {name, value} of the available fan speeds,
the main driver being the need to simplify homeassistant integrations for different devices.

For viomi vacuums this is currently straightforward and hard-coded, but we do special handling
for rockrobo devices (based on info() responses):

* If the query succeeds, newer firmware versions (3.5.7+) of v1 are handled as a special case,
  otherwise the new-style [100-105] mapping is returned.

* If the query fails, we fall back to rockrobo v1's percentage-based mapping.
  This happens, e.g., when the vacuum has no cloud connectivity.

Related to #523

Related downstream issues
https://github.com/home-assistant/core/pull/32821
https://github.com/home-assistant/core/issues/31268
https://github.com/home-assistant/core/issues/27268
---
 miio/vacuum.py      | 63 ++++++++++++++++++++++++++++++++++++++++++++-
 miio/viomivacuum.py |  6 +++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/miio/vacuum.py b/miio/vacuum.py
index 0ee49cf51..91412872b 100644
--- a/miio/vacuum.py
+++ b/miio/vacuum.py
@@ -6,7 +6,7 @@
 import os
 import pathlib
 import time
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import click
 import pytz
@@ -46,6 +46,24 @@ class Consumable(enum.Enum):
     SensorDirty = "sensor_dirty_time"
 
 
+class FanspeedV1(enum.Enum):
+    Silent = 38
+    Standard = 60
+    Medium = 77
+    Turbo = 90
+
+
+class FanspeedV2(enum.Enum):
+    Silent = 101
+    Standard = 102
+    Medium = 103
+    Turbo = 104
+    Gentle = 105
+
+
+ROCKROBO_V1 = "rockrobo.vacuum.v1"
+
+
 class Vacuum(Device):
     """Main class representing the vacuum."""
 
@@ -54,6 +72,8 @@ def __init__(
     ) -> None:
         super().__init__(ip, token, start_id, debug)
         self.manual_seqnum = -1
+        self.model = None
+        self._fanspeeds = FanspeedV1
 
     @command()
     def start(self):
@@ -416,6 +436,47 @@ def fan_speed(self):
         """Return fan speed."""
         return self.send("get_custom_mode")[0]
 
+    def _autodetect_model(self):
+        """Detect the model of the vacuum.
+
+        For the moment this is used only for the fanspeeds,
+        but that could be extended to cover other supported features."""
+        # cloud-blocked vacuums will not return proper payloads
+        try:
+            info = self.info()
+            self.model = info.model
+        except TypeError:
+            self._fanspeeds = FanspeedV1
+            self.model = ROCKROBO_V1
+            _LOGGER.debug("Unable to query model, falling back to %s", self._fanspeeds)
+            return
+
+        _LOGGER.info("model: %s", self.model)
+
+        if info.model == ROCKROBO_V1:
+            _LOGGER.info("Got robov1, checking for firmware version")
+            fw_version = info.firmware_version
+            version, build = fw_version.split("_")
+            version = tuple(map(int, version.split(".")))
+            if version >= (3, 5, 7):
+                self._fanspeeds = FanspeedV2
+            else:
+                self._fanspeeds = FanspeedV1
+        else:
+            self._fanspeeds = FanspeedV2
+
+        _LOGGER.debug(
+            "Using new fanspeed mapping %s for %s", self._fanspeeds, info.model
+        )
+
+    @command()
+    def supported_fanspeeds(self) -> Dict[str, int]:
+        """Return dictionary containing supported fanspeeds."""
+        if self.model is None:
+            self._autodetect_model()
+
+        return {x.name: x.value for x in list(self._fanspeeds)}
+
     @command()
     def sound_info(self):
         """Get voice settings."""
diff --git a/miio/viomivacuum.py b/miio/viomivacuum.py
index 67d370380..75ea871aa 100644
--- a/miio/viomivacuum.py
+++ b/miio/viomivacuum.py
@@ -3,6 +3,7 @@
 from collections import defaultdict
 from datetime import timedelta
 from enum import Enum
+from typing import Dict
 
 import click
 
@@ -338,3 +339,8 @@ def led(self, state: ViomiLedState):
     def carpet_mode(self, mode: ViomiCarpetTurbo):
         """Set the carpet mode."""
         return self.send("set_carpetturbo", [mode.value])
+
+    @command()
+    def supported_fanspeeds(self) -> Dict[str, int]:
+        """Return dictionary containing supported fanspeeds."""
+        return {x.name: x.value for x in list(ViomiVacuumSpeed)}