From 7421156dfc5d4a17384a708c623c28ad73a737a6 Mon Sep 17 00:00:00 2001
From: Sebastian Muszynski <basti@linkt.de>
Date: Sun, 24 Mar 2019 20:15:29 +0100
Subject: [PATCH] Add support for the power socket of the Xiaomi AC Partner V3
 (#22205)

---
 .../components/xiaomi_miio/switch.py          | 80 +++++++++++++++----
 1 file changed, 65 insertions(+), 15 deletions(-)

diff --git a/homeassistant/components/xiaomi_miio/switch.py b/homeassistant/components/xiaomi_miio/switch.py
index f330030922d..d1acce02e47 100644
--- a/homeassistant/components/xiaomi_miio/switch.py
+++ b/homeassistant/components/xiaomi_miio/switch.py
@@ -36,6 +36,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
          'chuangmi.plug.v2',
          'chuangmi.plug.v3',
          'chuangmi.plug.hmi205',
+         'lumi.acpartner.v3',
          ]),
 })
 
@@ -150,6 +151,13 @@ async def async_setup_platform(hass, config, async_add_entities,
         device = XiaomiPlugGenericSwitch(name, plug, model, unique_id)
         devices.append(device)
         hass.data[DATA_KEY][host] = device
+    elif model in ['lumi.acpartner.v3']:
+        from miio import AirConditioningCompanionV3
+        plug = AirConditioningCompanionV3(host, token)
+        device = XiaomiAirConditioningCompanionSwitch(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 '
@@ -294,9 +302,7 @@ class XiaomiPlugGenericSwitch(SwitchDevice):
 
             self._available = True
             self._state = state.is_on
-            self._state_attrs.update({
-                ATTR_TEMPERATURE: state.temperature
-            })
+            self._state_attrs[ATTR_TEMPERATURE] = state.temperature
 
         except DeviceException as ex:
             self._available = False
@@ -342,9 +348,7 @@ class XiaomiPowerStripSwitch(XiaomiPlugGenericSwitch):
         else:
             self._device_features = FEATURE_FLAGS_POWER_STRIP_V1
 
-        self._state_attrs.update({
-            ATTR_LOAD_POWER: None,
-        })
+        self._state_attrs[ATTR_LOAD_POWER] = None
 
         if self._device_features & FEATURE_SET_POWER_MODE == 1:
             self._state_attrs[ATTR_POWER_MODE] = None
@@ -418,13 +422,9 @@ class ChuangMiPlugSwitch(XiaomiPlugGenericSwitch):
 
         if self._model == MODEL_PLUG_V3:
             self._device_features = FEATURE_FLAGS_PLUG_V3
-            self._state_attrs.update({
-                ATTR_WIFI_LED: None,
-            })
+            self._state_attrs[ATTR_WIFI_LED] = None
             if self._channel_usb is False:
-                self._state_attrs.update({
-                    ATTR_LOAD_POWER: None,
-                })
+                self._state_attrs[ATTR_LOAD_POWER] = None
 
     async def async_turn_on(self, **kwargs):
         """Turn a channel on."""
@@ -471,9 +471,7 @@ class ChuangMiPlugSwitch(XiaomiPlugGenericSwitch):
             else:
                 self._state = state.is_on
 
-            self._state_attrs.update({
-                ATTR_TEMPERATURE: state.temperature
-            })
+            self._state_attrs[ATTR_TEMPERATURE] = state.temperature
 
             if state.wifi_led:
                 self._state_attrs[ATTR_WIFI_LED] = state.wifi_led
@@ -484,3 +482,55 @@ class ChuangMiPlugSwitch(XiaomiPlugGenericSwitch):
         except DeviceException as ex:
             self._available = False
             _LOGGER.error("Got exception while fetching the state: %s", ex)
+
+
+class XiaomiAirConditioningCompanionSwitch(XiaomiPlugGenericSwitch):
+    """Representation of a Xiaomi AirConditioning Companion."""
+
+    def __init__(self, name, plug, model, unique_id):
+        """Initialize the acpartner switch."""
+        super().__init__(name, plug, model, unique_id)
+
+        self._state_attrs.update({
+            ATTR_TEMPERATURE: None,
+            ATTR_LOAD_POWER: None,
+        })
+
+    async def async_turn_on(self, **kwargs):
+        """Turn the socket on."""
+        result = await self._try_command(
+            "Turning the socket on failed.", self._plug.socket_on)
+
+        if result:
+            self._state = True
+            self._skip_update = True
+
+    async def async_turn_off(self, **kwargs):
+        """Turn the socket off."""
+        result = await self._try_command(
+            "Turning the socket off failed.", self._plug.socket_off)
+
+        if result:
+            self._state = False
+            self._skip_update = True
+
+    async def async_update(self):
+        """Fetch state from the device."""
+        from miio import DeviceException
+
+        # On state change the device doesn't provide the new state immediately.
+        if self._skip_update:
+            self._skip_update = False
+            return
+
+        try:
+            state = await self.hass.async_add_executor_job(self._plug.status)
+            _LOGGER.debug("Got new state: %s", state)
+
+            self._available = True
+            self._state = state.power_socket == 'on'
+            self._state_attrs[ATTR_LOAD_POWER] = state.load_power
+
+        except DeviceException as ex:
+            self._available = False
+            _LOGGER.error("Got exception while fetching the state: %s", ex)