From 5f748cbd585fd977e454fc8ba263791bf04f4b0f Mon Sep 17 00:00:00 2001
From: Teemu Rytilahti <tpr@iki.fi>
Date: Sun, 16 Jan 2022 17:25:02 +0100
Subject: [PATCH] Improve miotdevice mappings handling

* Introduce _mappings containing (model => mapping) to allow easier support
  for different device models
* Fallback to first _mappings entry when encountering a model without mapping
* Use `mapping` for backwards compatibility for existing code

* Convert FanMiot to use the mappings dict, deprecate FanP9, FanP10, FanP11
---
 miio/fan_miot.py              | 20 +++++-------------
 miio/miot_device.py           | 38 +++++++++++++++++++++++++++++------
 miio/tests/test_miotdevice.py | 23 +++++++++++++++++++++
 3 files changed, 60 insertions(+), 21 deletions(-)

diff --git a/miio/fan_miot.py b/miio/fan_miot.py
index 86558edb4..b12ae1560 100644
--- a/miio/fan_miot.py
+++ b/miio/fan_miot.py
@@ -6,6 +6,7 @@
 from .click_common import EnumType, command, format_output
 from .fan_common import FanException, MoveDirection, OperationMode
 from .miot_device import DeviceStatus, MiotDevice
+from .utils import deprecated
 
 MODEL_FAN_P9 = "dmaker.fan.p9"
 MODEL_FAN_P10 = "dmaker.fan.p10"
@@ -252,21 +253,7 @@ def child_lock(self) -> bool:
 
 
 class FanMiot(MiotDevice):
-    mapping = MIOT_MAPPING[MODEL_FAN_P10]
-
-    def __init__(
-        self,
-        ip: str = None,
-        token: str = None,
-        start_id: int = 0,
-        debug: int = 0,
-        lazy_discover: bool = True,
-        model: str = MODEL_FAN_P10,
-    ) -> None:
-        if model not in MIOT_MAPPING:
-            raise FanException("Invalid FanMiot model: %s" % model)
-
-        super().__init__(ip, token, start_id, debug, lazy_discover, model=model)
+    _mappings = MIOT_MAPPING
 
     @command(
         default_output=format_output(
@@ -406,14 +393,17 @@ def set_rotate(self, direction: MoveDirection):
         return self.set_property("set_move", value)
 
 
+@deprecated("Use FanMiot")
 class FanP9(FanMiot):
     mapping = MIOT_MAPPING[MODEL_FAN_P9]
 
 
+@deprecated("Use FanMiot")
 class FanP10(FanMiot):
     mapping = MIOT_MAPPING[MODEL_FAN_P10]
 
 
+@deprecated("Use FanMiot")
 class FanP11(FanMiot):
     mapping = MIOT_MAPPING[MODEL_FAN_P11]
 
diff --git a/miio/miot_device.py b/miio/miot_device.py
index d53557454..89559d600 100644
--- a/miio/miot_device.py
+++ b/miio/miot_device.py
@@ -28,9 +28,13 @@ def _str2bool(x):
 
 
 class MiotDevice(Device):
-    """Main class representing a MIoT device."""
+    """Main class representing a MIoT device.
+
+    The implementors should extend the
+    """
 
     mapping: MiotMapping
+    _mappings: Dict[str, MiotMapping] = {}
 
     def __init__(
         self,
@@ -49,7 +53,7 @@ def __init__(
             ip, token, start_id, debug, lazy_discover, timeout, model=model
         )
 
-        if mapping is None and not hasattr(self, "mapping"):
+        if mapping is None and not hasattr(self, "mapping") and not self._mappings:
             _LOGGER.warning("Neither the class nor the parameter defines the mapping")
 
         if mapping is not None:
@@ -59,9 +63,8 @@ def get_properties_for_mapping(self, *, max_properties=15) -> list:
         """Retrieve raw properties based on mapping."""
 
         # We send property key in "did" because it's sent back via response and we can identify the property.
-        properties = [
-            {"did": k, **v} for k, v in self.mapping.items() if "aiid" not in v
-        ]
+        mapping = self._get_mapping()
+        properties = [{"did": k, **v} for k, v in mapping.items() if "aiid" not in v]
 
         return self.get_properties(
             properties, property_getter="get_properties", max_properties=max_properties
@@ -141,7 +144,30 @@ def set_property_by(
 
     def set_property(self, property_key: str, value):
         """Sets property value using the existing mapping."""
+        mapping = self._get_mapping()
         return self.send(
             "set_properties",
-            [{"did": property_key, **self.mapping[property_key], "value": value}],
+            [{"did": property_key, **mapping[property_key], "value": value}],
         )
+
+    def _get_mapping(self) -> MiotMapping:
+        """Return the protocol mapping to use.
+
+        The logic is as follows:
+        1. Use device model as key to lookup _mappings for the mapping
+        2. If no match is found, but _mappings is defined, use the first item
+        3. Fallback to class-defined `mapping` for backwards compat
+        """
+        if not self._mappings:
+            return self.mapping
+
+        mapping = self._mappings.get(self.model)
+        if mapping is not None:
+            return mapping
+
+        first_model, first_mapping = list(self._mappings.items())[0]
+        _LOGGER.warning(
+            "Unable to find mapping for %s, falling back to %s", self.model, first_model
+        )
+
+        return first_mapping
diff --git a/miio/tests/test_miotdevice.py b/miio/tests/test_miotdevice.py
index 429e85d40..7bfe8ddac 100644
--- a/miio/tests/test_miotdevice.py
+++ b/miio/tests/test_miotdevice.py
@@ -90,3 +90,26 @@ def test_call_action_by(dev):
             "in": params,
         },
     )
+
+
+@pytest.mark.parametrize(
+    "model,expected_mapping,expected_log",
+    [
+        ("some_model", {"x": {"y": 1}}, ""),
+        ("unknown_model", {"x": {"y": 1}}, "Unable to find mapping"),
+    ],
+)
+def test_get_mapping(dev, caplog, model, expected_mapping, expected_log):
+    """Test _get_mapping logic for fallbacks."""
+    dev._mappings["some_model"] = {"x": {"y": 1}}
+    dev._model = model
+    assert dev._get_mapping() == expected_mapping
+
+    assert expected_log in caplog.text
+
+
+def test_get_mapping_backwards_compat(dev):
+    """Test that the backwards compat works."""
+    # as dev is mocked on module level, need to empty manually
+    dev._mappings = {}
+    assert dev._get_mapping() == {}