Powerstrip support improved

miio/powerstrip.py

@@ -2,11 +2,15 @@
 import enum
 from typing import Dict, Any, Optional
 from collections import defaultdict
-from .device import Device
+from .device import Device, DeviceException
 
 _LOGGER = logging.getLogger(__name__)
 
 
+class PowerStripException(DeviceException):
+    pass
+
+
 class PowerMode(enum.Enum):
     Eco = 'green'
     Normal = 'normal'
@@ -16,10 +20,13 @@ class PowerStripStatus:
     """Container for status reports from the power strip."""
 
     def __init__(self, data: Dict[str, Any]) -> None:
-        # Device model: qmi.powerstrip.v1, zimi.powerstrip.v2
-        #
-        # {'power': 'on', 'temperature': 48.11,
-        # 'current': 0.06, 'mode': 'green'}
+        """
+        Supported device models: qmi.powerstrip.v1, zimi.powerstrip.v2
+
+        Response of a Power Strip 2 (zimi.powerstrip.v2):
+        {'power','on', 'temperature': 48.7, 'current': 0.05, 'mode': None,
+         'power_consume_rate': 4.09, 'wifi_led': 'on', 'power_price': 49}
+        """
         self.data = data
 
     @property
@@ -58,14 +65,33 @@ def mode(self) -> Optional[PowerMode]:
             return PowerMode(self.data["mode"])
         return None
 
+    @property
+    def wifi_led(self) -> bool:
+        """True if the wifi led is turned on."""
+        return self.data["wifi_led"] == "on"
+
+    @property
+    def power_price(self) -> Optional[int]:
+        """The stored power price, if available."""
+        if self.data["power_price"] is not None:
+            return self.data["power_price"]
+        return None
+
     def __repr__(self) -> str:
-        s = "<PowerStripStatus power=%s, temperature=%s, " \
-            "load_power=%s, current=%s, mode=%s>" % \
+        s = "<PowerStripStatus power=%s, " \
+            "temperature=%s, " \
+            "load_power=%s, " \
+            "current=%s, " \
+            "mode=%s, " \
+            "wifi_led=%s, " \
+            "power_price=%s>" % \
             (self.power,
              self.temperature,
              self.load_power,
              self.current,
-             self.mode)
+             self.mode,
+             self.wifi_led,
+             self.power_price)
         return s
 
 
@@ -75,7 +101,7 @@ class PowerStrip(Device):
     def status(self) -> PowerStripStatus:
         """Retrieve properties."""
         properties = ['power', 'temperature', 'current', 'mode',
-                      'power_consume_rate']
+                      'power_consume_rate', 'wifi_led', 'power_price', ]
         values = self.send(
             "get_prop",
             properties
@@ -101,7 +127,21 @@ def off(self):
         return self.send("set_power", ["off"])
 
     def set_power_mode(self, mode: PowerMode):
-        """Set mode."""
+        """Set the power mode."""
 
         # green, normal
         return self.send("set_power_mode", [mode.value])
+
+    def set_wifi_led(self, led: bool):
+        """Set the wifi led on/off."""
+        if led:
+            return self.send("set_wifi_led", ["on"])
+        else:
+            return self.send("set_wifi_led", ["off"])
+
+    def set_power_price(self, price: int):
+        """Set the power price."""
+        if price < 0:
+            raise PowerStripException("Invalid power price: %s" % price)
+
+        return self.send("set_power_price", [price])

miio/tests/test_powerstrip.py

@@ -1,6 +1,6 @@
 from unittest import TestCase
 from miio import PowerStrip
-from miio.powerstrip import PowerMode
+from miio.powerstrip import PowerMode, PowerStripException
 from .dummies import DummyDevice
 import pytest
 
@@ -13,11 +13,15 @@ def __init__(self, *args, **kwargs):
             'temperature': 32.5,
             'current': 25.5,
             'power_consume_rate': 12.5,
+            'wifi_led': 'off',
+            'power_price': 49,
         }
         self.return_values = {
             'get_prop': self._get_state,
             'set_power': lambda x: self._set_state("power", x),
             'set_power_mode': lambda x: self._set_state("mode", x),
+            'set_wifi_led': lambda x: self._set_state("wifi_led", x),
+            'set_power_price': lambda x: self._set_state("power_price", x),
         }
         super().__init__(args, kwargs)
 
@@ -87,3 +91,40 @@ def mode():
         assert mode() == PowerMode.Eco
         self.device.set_power_mode(PowerMode.Normal)
         assert mode() == PowerMode.Normal
+
+    def test_status_without_mode(self):
+        self.device._reset_state()
+
+        # The Power Strip  2 doesn't support power modes
+        self.device.state["mode"] = None
+        assert self.state().mode is None
+
+    def test_set_wifi_led(self):
+        def wifi_led():
+            return self.device.status().wifi_led
+
+        self.device.set_wifi_led(True)
+        assert wifi_led() is True
+
+        self.device.set_wifi_led(False)
+        assert wifi_led() is False
+
+    def test_set_power_price(self):
+        def power_price():
+            return self.device.status().power_price
+
+        self.device.set_power_price(0)
+        assert power_price() == 0
+        self.device.set_power_price(1)
+        assert power_price() == 1
+        self.device.set_power_price(2)
+        assert power_price() == 2
+
+        with pytest.raises(PowerStripException):
+            self.device.set_power_price(-1)
+
+    def test_status_without_power_price(self):
+        self.device._reset_state()
+
+        self.device.state["power_price"] = None
+        assert self.state().power_price is None