Xiaomi MiIO Switch: Power Strip support improved (#12917)

* Xiaomi MiIO Switch: Power Strip support improved.

* New service descriptions added.

* Make hound happy.

* Pylint fixed.

* Use Async / await syntax.

* Missed method fixed.

* Make hound happy.

* Don't abuse the system property supported_features anymore.

* Check the correct method.

* Refactoring.

* Make hound happy.

* pythion-miio version bumped.

* Clean-up.

* Unique id added.

* Filter service calls.
Device unavailable handling improved.

homeassistant/components/switch/services.yaml

@@ -30,3 +30,34 @@ mysensors_send_ir_code:
     V_IR_SEND:
       description: IR code to send.
       example: '0xC284'
+
+xiaomi_miio_set_wifi_led_on:
+  description: Turn the wifi led on.
+  fields:
+    entity_id:
+      description: Name of the xiaomi miio entity.
+      example: 'switch.xiaomi_miio_device'
+xiaomi_miio_set_wifi_led_off:
+  description: Turn the wifi led off.
+  fields:
+    entity_id:
+      description: Name of the xiaomi miio entity.
+      example: 'switch.xiaomi_miio_device'
+xiaomi_miio_set_power_price:
+  description: Set the power price.
+  fields:
+    entity_id:
+      description: Name of the xiaomi miio entity.
+      example: 'switch.xiaomi_miio_device'
+    mode:
+      description: Power price, between 0 and 999.
+      example: 31
+xiaomi_miio_set_power_mode:
+  description: Set the power mode.
+  fields:
+    entity_id:
+      description: Name of the xiaomi miio entity.
+      example: 'switch.xiaomi_miio_device'
+    mode:
+      description: Power mode, valid values are 'normal' and 'green'.
+      example: 'green'

homeassistant/components/switch/xiaomi_miio.py

@@ -11,15 +11,19 @@ import logging
 import voluptuous as vol
 
 import homeassistant.helpers.config_validation as cv
-from homeassistant.components.switch import (SwitchDevice, PLATFORM_SCHEMA, )
+from homeassistant.components.switch import (SwitchDevice, PLATFORM_SCHEMA,
-from homeassistant.const import (CONF_NAME, CONF_HOST, CONF_TOKEN, )
+                                             DOMAIN, )
+from homeassistant.const import (CONF_NAME, CONF_HOST, CONF_TOKEN,
+                                 ATTR_ENTITY_ID, )
 from homeassistant.exceptions import PlatformNotReady
 
 _LOGGER = logging.getLogger(__name__)
 
 DEFAULT_NAME = 'Xiaomi Miio Switch'
+DATA_KEY = 'switch.xiaomi_miio'
 
 CONF_MODEL = 'model'
+MODEL_POWER_STRIP_V2 = 'zimi.powerstrip.v2'
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_HOST): cv.string,
@@ -39,14 +43,63 @@ ATTR_POWER = 'power'
 ATTR_TEMPERATURE = 'temperature'
 ATTR_LOAD_POWER = 'load_power'
 ATTR_MODEL = 'model'
+ATTR_MODE = 'mode'
+ATTR_POWER_MODE = 'power_mode'
+ATTR_WIFI_LED = 'wifi_led'
+ATTR_POWER_PRICE = 'power_price'
+ATTR_PRICE = 'price'
+
 SUCCESS = ['ok']
 
+SUPPORT_SET_POWER_MODE = 1
+SUPPORT_SET_WIFI_LED = 2
+SUPPORT_SET_POWER_PRICE = 4
+
+ADDITIONAL_SUPPORT_FLAGS_GENERIC = 0
+
+ADDITIONAL_SUPPORT_FLAGS_POWER_STRIP_V1 = (SUPPORT_SET_POWER_MODE |
+                                           SUPPORT_SET_WIFI_LED |
+                                           SUPPORT_SET_POWER_PRICE)
+
+ADDITIONAL_SUPPORT_FLAGS_POWER_STRIP_V2 = (SUPPORT_SET_WIFI_LED |
+                                           SUPPORT_SET_POWER_PRICE)
+
+SERVICE_SET_WIFI_LED_ON = 'xiaomi_miio_set_wifi_led_on'
+SERVICE_SET_WIFI_LED_OFF = 'xiaomi_miio_set_wifi_led_off'
+SERVICE_SET_POWER_MODE = 'xiaomi_miio_set_power_mode'
+SERVICE_SET_POWER_PRICE = 'xiaomi_miio_set_power_price'
+
+SERVICE_SCHEMA = vol.Schema({
+    vol.Optional(ATTR_ENTITY_ID): cv.entity_ids,
+})
+
+SERVICE_SCHEMA_POWER_MODE = SERVICE_SCHEMA.extend({
+    vol.Required(ATTR_MODE): vol.All(vol.In(['green', 'normal'])),
+})
+
+SERVICE_SCHEMA_POWER_PRICE = SERVICE_SCHEMA.extend({
+    vol.Required(ATTR_PRICE): vol.All(vol.Coerce(float), vol.Range(min=0))
+})
+
+SERVICE_TO_METHOD = {
+    SERVICE_SET_WIFI_LED_ON: {'method': 'async_set_wifi_led_on'},
+    SERVICE_SET_WIFI_LED_OFF: {'method': 'async_set_wifi_led_off'},
+    SERVICE_SET_POWER_MODE: {
+        'method': 'async_set_power_mode',
+        'schema': SERVICE_SCHEMA_POWER_MODE},
+    SERVICE_SET_POWER_PRICE: {
+        'method': 'async_set_power_price',
+        'schema': SERVICE_SCHEMA_POWER_PRICE},
+}
+
 
 # pylint: disable=unused-argument
-@asyncio.coroutine
+async def async_setup_platform(hass, config, async_add_devices,
-def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
+                               discovery_info=None):
     """Set up the switch from config."""
     from miio import Device, DeviceException
+    if DATA_KEY not in hass.data:
+        hass.data[DATA_KEY] = {}
 
     host = config.get(CONF_HOST)
     name = config.get(CONF_NAME)
@@ -56,12 +109,14 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
     _LOGGER.info("Initializing with host %s (token %s...)", host, token[:5])
 
     devices = []
+    unique_id = None
 
     if model is None:
         try:
             miio_device = Device(host, token)
             device_info = miio_device.info()
             model = device_info.model
+            unique_id = "{}-{}".format(model, device_info.mac_address)
             _LOGGER.info("%s %s %s detected",
                          model,
                          device_info.firmware_version,
@@ -77,21 +132,24 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
         # A switch device per channel will be created.
         for channel_usb in [True, False]:
             device = ChuangMiPlugV1Switch(
-                name, plug, model, channel_usb)
+                name, plug, model, unique_id, channel_usb)
             devices.append(device)
+            hass.data[DATA_KEY][host] = device
 
     elif model in ['qmi.powerstrip.v1',
                    'zimi.powerstrip.v2']:
         from miio import PowerStrip
         plug = PowerStrip(host, token)
-        device = XiaomiPowerStripSwitch(name, plug, model)
+        device = XiaomiPowerStripSwitch(name, plug, model, unique_id)
         devices.append(device)
+        hass.data[DATA_KEY][host] = device
     elif model in ['chuangmi.plug.m1',
                    'chuangmi.plug.v2']:
         from miio import Plug
         plug = Plug(host, token)
-        device = XiaomiPlugGenericSwitch(name, plug, model)
+        device = XiaomiPlugGenericSwitch(name, plug, model, unique_id)
         devices.append(device)
+        hass.data[DATA_KEY][host] = device
     else:
         _LOGGER.error(
             'Unsupported device found! Please create an issue at '
@@ -101,22 +159,52 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
 
     async_add_devices(devices, update_before_add=True)
 
+    async def async_service_handler(service):
+        """Map services to methods on XiaomiPlugGenericSwitch."""
+        method = SERVICE_TO_METHOD.get(service.service)
+        params = {key: value for key, value in service.data.items()
+                  if key != ATTR_ENTITY_ID}
+        entity_ids = service.data.get(ATTR_ENTITY_ID)
+        if entity_ids:
+            devices = [device for device in hass.data[DATA_KEY].values() if
+                       device.entity_id in entity_ids]
+        else:
+            devices = hass.data[DATA_KEY].values()
+
+        update_tasks = []
+        for device in devices:
+            if not hasattr(device, method['method']):
+                continue
+            await getattr(device, method['method'])(**params)
+            update_tasks.append(device.async_update_ha_state(True))
+
+        if update_tasks:
+            await asyncio.wait(update_tasks, loop=hass.loop)
+
+    for plug_service in SERVICE_TO_METHOD:
+        schema = SERVICE_TO_METHOD[plug_service].get('schema', SERVICE_SCHEMA)
+        hass.services.async_register(
+            DOMAIN, plug_service, async_service_handler, schema=schema)
+
 
 class XiaomiPlugGenericSwitch(SwitchDevice):
     """Representation of a Xiaomi Plug Generic."""
 
-    def __init__(self, name, plug, model):
+    def __init__(self, name, plug, model, unique_id):
         """Initialize the plug switch."""
         self._name = name
-        self._icon = 'mdi:power-socket'
-        self._model = model
-
         self._plug = plug
+        self._model = model
+        self._unique_id = unique_id
+
+        self._icon = 'mdi:power-socket'
+        self._available = False
         self._state = None
         self._state_attrs = {
             ATTR_TEMPERATURE: None,
             ATTR_MODEL: self._model,
         }
+        self._additional_supported_features = ADDITIONAL_SUPPORT_FLAGS_GENERIC
         self._skip_update = False
 
     @property
@@ -124,6 +212,11 @@ class XiaomiPlugGenericSwitch(SwitchDevice):
         """Poll the plug."""
         return True
 
+    @property
+    def unique_id(self):
+        """Return an unique ID."""
+        return self._unique_id
+
     @property
     def name(self):
         """Return the name of the device if any."""
@@ -137,7 +230,7 @@ class XiaomiPlugGenericSwitch(SwitchDevice):
     @property
     def available(self):
         """Return true when state is known."""
-        return self._state is not None
+        return self._available
 
     @property
     def device_state_attributes(self):
@@ -149,12 +242,11 @@ class XiaomiPlugGenericSwitch(SwitchDevice):
         """Return true if switch is on."""
         return self._state
 
-    @asyncio.coroutine
+    async def _try_command(self, mask_error, func, *args, **kwargs):
-    def _try_command(self, mask_error, func, *args, **kwargs):
         """Call a plug command handling error messages."""
         from miio import DeviceException
         try:
-            result = yield from self.hass.async_add_job(
+            result = await self.hass.async_add_job(
                 partial(func, *args, **kwargs))
 
             _LOGGER.debug("Response received from plug: %s", result)
@@ -162,30 +254,28 @@ class XiaomiPlugGenericSwitch(SwitchDevice):
             return result == SUCCESS
         except DeviceException as exc:
             _LOGGER.error(mask_error, exc)
+            self._available = False
             return False
 
-    @asyncio.coroutine
+    async def async_turn_on(self, **kwargs):
-    def async_turn_on(self, **kwargs):
         """Turn the plug on."""
-        result = yield from self._try_command(
+        result = await self._try_command(
             "Turning the plug on failed.", self._plug.on)
 
         if result:
             self._state = True
             self._skip_update = True
 
-    @asyncio.coroutine
+    async def async_turn_off(self, **kwargs):
-    def async_turn_off(self, **kwargs):
         """Turn the plug off."""
-        result = yield from self._try_command(
+        result = await self._try_command(
             "Turning the plug off failed.", self._plug.off)
 
         if result:
             self._state = False
             self._skip_update = True
 
-    @asyncio.coroutine
+    async def async_update(self):
-    def async_update(self):
         """Fetch state from the device."""
         from miio import DeviceException
 
@@ -195,34 +285,75 @@ class XiaomiPlugGenericSwitch(SwitchDevice):
             return
 
         try:
-            state = yield from self.hass.async_add_job(self._plug.status)
+            state = await self.hass.async_add_job(self._plug.status)
             _LOGGER.debug("Got new state: %s", state)
 
+            self._available = True
             self._state = state.is_on
             self._state_attrs.update({
                 ATTR_TEMPERATURE: state.temperature
             })
 
         except DeviceException as ex:
-            self._state = None
+            self._available = False
             _LOGGER.error("Got exception while fetching the state: %s", ex)
 
+    async def async_set_wifi_led_on(self):
+        """Turn the wifi led on."""
+        if self._additional_supported_features & SUPPORT_SET_WIFI_LED == 0:
+            return
 
-class XiaomiPowerStripSwitch(XiaomiPlugGenericSwitch, SwitchDevice):
+        await self._try_command(
+            "Turning the wifi led on failed.",
+            self._plug.set_wifi_led, True)
+
+    async def async_set_wifi_led_off(self):
+        """Turn the wifi led on."""
+        if self._additional_supported_features & SUPPORT_SET_WIFI_LED == 0:
+            return
+
+        await self._try_command(
+            "Turning the wifi led off failed.",
+            self._plug.set_wifi_led, False)
+
+    async def async_set_power_price(self, price: int):
+        """Set the power price."""
+        if self._additional_supported_features & SUPPORT_SET_POWER_PRICE == 0:
+            return
+
+        await self._try_command(
+            "Setting the power price of the power strip failed.",
+            self._plug.set_power_price, price)
+
+
+class XiaomiPowerStripSwitch(XiaomiPlugGenericSwitch):
     """Representation of a Xiaomi Power Strip."""
 
-    def __init__(self, name, plug, model):
+    def __init__(self, name, plug, model, unique_id):
         """Initialize the plug switch."""
-        XiaomiPlugGenericSwitch.__init__(self, name, plug, model)
+        XiaomiPlugGenericSwitch.__init__(self, name, plug, model, unique_id)
 
-        self._state_attrs = {
+        if self._model == MODEL_POWER_STRIP_V2:
-            ATTR_TEMPERATURE: None,
+            self._additional_supported_features = \
+                ADDITIONAL_SUPPORT_FLAGS_POWER_STRIP_V2
+        else:
+            self._additional_supported_features = \
+                ADDITIONAL_SUPPORT_FLAGS_POWER_STRIP_V1
+
+        self._state_attrs.update({
             ATTR_LOAD_POWER: None,
-            ATTR_MODEL: self._model,
+        })
-        }
 
-    @asyncio.coroutine
+        if self._additional_supported_features & SUPPORT_SET_POWER_MODE == 1:
-    def async_update(self):
+            self._state_attrs[ATTR_POWER_MODE] = None
+
+        if self._additional_supported_features & SUPPORT_SET_WIFI_LED == 1:
+            self._state_attrs[ATTR_WIFI_LED] = None
+
+        if self._additional_supported_features & SUPPORT_SET_POWER_PRICE == 1:
+            self._state_attrs[ATTR_POWER_PRICE] = None
+
+    async def async_update(self):
         """Fetch state from the device."""
         from miio import DeviceException
 
@@ -232,60 +363,84 @@ class XiaomiPowerStripSwitch(XiaomiPlugGenericSwitch, SwitchDevice):
             return
 
         try:
-            state = yield from self.hass.async_add_job(self._plug.status)
+            state = await self.hass.async_add_job(self._plug.status)
             _LOGGER.debug("Got new state: %s", state)
 
+            self._available = True
             self._state = state.is_on
             self._state_attrs.update({
                 ATTR_TEMPERATURE: state.temperature,
-                ATTR_LOAD_POWER: state.load_power
+                ATTR_LOAD_POWER: state.load_power,
             })
 
+            if self._additional_supported_features & \
+                    SUPPORT_SET_POWER_MODE == 1 and state.mode:
+                self._state_attrs[ATTR_POWER_MODE] = state.mode.value
+
+            if self._additional_supported_features & \
+                    SUPPORT_SET_WIFI_LED == 1 and state.wifi_led:
+                self._state_attrs[ATTR_WIFI_LED] = state.wifi_led
+
+            if self._additional_supported_features & \
+                    SUPPORT_SET_POWER_PRICE == 1 and state.power_price:
+                self._state_attrs[ATTR_POWER_PRICE] = state.power_price
+
         except DeviceException as ex:
-            self._state = None
+            self._available = False
             _LOGGER.error("Got exception while fetching the state: %s", ex)
 
+    async def async_set_power_mode(self, mode: str):
+        """Set the power mode."""
+        if self._additional_supported_features & SUPPORT_SET_POWER_MODE == 0:
+            return
 
-class ChuangMiPlugV1Switch(XiaomiPlugGenericSwitch, SwitchDevice):
+        from miio.powerstrip import PowerMode
+
+        await self._try_command(
+            "Setting the power mode of the power strip failed.",
+            self._plug.set_power_mode, PowerMode(mode))
+
+
+class ChuangMiPlugV1Switch(XiaomiPlugGenericSwitch):
     """Representation of a Chuang Mi Plug V1."""
 
-    def __init__(self, name, plug, model, channel_usb):
+    def __init__(self, name, plug, model, unique_id, channel_usb):
         """Initialize the plug switch."""
         name = '{} USB'.format(name) if channel_usb else name
 
-        XiaomiPlugGenericSwitch.__init__(self, name, plug, model)
+        if unique_id is not None and channel_usb:
+            unique_id = "{}-{}".format(unique_id, 'usb')
+
+        XiaomiPlugGenericSwitch.__init__(self, name, plug, model, unique_id)
         self._channel_usb = channel_usb
 
-    @asyncio.coroutine
+    async def async_turn_on(self, **kwargs):
-    def async_turn_on(self, **kwargs):
         """Turn a channel on."""
         if self._channel_usb:
-            result = yield from self._try_command(
+            result = await self._try_command(
                 "Turning the plug on failed.", self._plug.usb_on)
         else:
-            result = yield from self._try_command(
+            result = await self._try_command(
                 "Turning the plug on failed.", self._plug.on)
 
         if result:
             self._state = True
             self._skip_update = True
 
-    @asyncio.coroutine
+    async def async_turn_off(self, **kwargs):
-    def async_turn_off(self, **kwargs):
         """Turn a channel off."""
         if self._channel_usb:
-            result = yield from self._try_command(
+            result = await self._try_command(
                 "Turning the plug on failed.", self._plug.usb_off)
         else:
-            result = yield from self._try_command(
+            result = await self._try_command(
                 "Turning the plug on failed.", self._plug.off)
 
         if result:
             self._state = False
             self._skip_update = True
 
-    @asyncio.coroutine
+    async def async_update(self):
-    def async_update(self):
         """Fetch state from the device."""
         from miio import DeviceException
 
@@ -295,9 +450,10 @@ class ChuangMiPlugV1Switch(XiaomiPlugGenericSwitch, SwitchDevice):
             return
 
         try:
-            state = yield from self.hass.async_add_job(self._plug.status)
+            state = await self.hass.async_add_job(self._plug.status)
             _LOGGER.debug("Got new state: %s", state)
 
+            self._available = True
             if self._channel_usb:
                 self._state = state.usb_power
             else:
@@ -308,5 +464,5 @@ class ChuangMiPlugV1Switch(XiaomiPlugGenericSwitch, SwitchDevice):
             })
 
         except DeviceException as ex:
-            self._state = None
+            self._available = False
             _LOGGER.error("Got exception while fetching the state: %s", ex)