From bb31de1de7320da94268a2277933ff9c323aceee Mon Sep 17 00:00:00 2001
From: b4dpxl <26406101+b4dpxl@users.noreply.github.com>
Date: Sat, 14 Nov 2020 13:59:41 +0000
Subject: [PATCH] Add support for Broadlink BG1 devices (#42314)

* Support for BG1 switches after config flow updates to core Broadlink component

* updates based on @felipediel feedback

* Update updater.py

* Update switch.py

* Update switch.py
---
 homeassistant/components/broadlink/const.py   |  2 +-
 homeassistant/components/broadlink/switch.py  | 47 +++++++++++++++++++
 homeassistant/components/broadlink/updater.py |  9 ++++
 3 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/broadlink/const.py b/homeassistant/components/broadlink/const.py
index adb5a437ae5..b10f7e74ba7 100644
--- a/homeassistant/components/broadlink/const.py
+++ b/homeassistant/components/broadlink/const.py
@@ -8,7 +8,7 @@ DOMAIN = "broadlink"
 DOMAINS_AND_TYPES = (
     (REMOTE_DOMAIN, ("RM2", "RM4")),
     (SENSOR_DOMAIN, ("A1", "RM2", "RM4")),
-    (SWITCH_DOMAIN, ("MP1", "RM2", "RM4", "SP1", "SP2", "SP4", "SP4B")),
+    (SWITCH_DOMAIN, ("BG1", "MP1", "RM2", "RM4", "SP1", "SP2", "SP4", "SP4B")),
 )
 
 DEFAULT_PORT = 80
diff --git a/homeassistant/components/broadlink/switch.py b/homeassistant/components/broadlink/switch.py
index 644255d7d17..b4cd43ac493 100644
--- a/homeassistant/components/broadlink/switch.py
+++ b/homeassistant/components/broadlink/switch.py
@@ -1,5 +1,6 @@
 """Support for Broadlink switches."""
 from abc import ABC, abstractmethod
+from functools import partial
 import logging
 
 from broadlink.exceptions import BroadlinkException
@@ -124,6 +125,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     elif device.api.type in {"SP4", "SP4B"}:
         switches = [BroadlinkSP4Switch(device)]
 
+    elif device.api.type == "BG1":
+        switches = [BroadlinkBG1Slot(device, slot) for slot in range(1, 3)]
+
     elif device.api.type == "MP1":
         switches = [BroadlinkMP1Slot(device, slot) for slot in range(1, 5)]
 
@@ -360,3 +364,46 @@ class BroadlinkMP1Slot(BroadlinkSwitch):
             _LOGGER.error("Failed to send packet: %s", err)
             return False
         return True
+
+
+class BroadlinkBG1Slot(BroadlinkSwitch):
+    """Representation of a Broadlink BG1 slot."""
+
+    def __init__(self, device, slot):
+        """Initialize the switch."""
+        super().__init__(device, 1, 0)
+        self._slot = slot
+        self._state = self._coordinator.data[f"pwr{slot}"]
+        self._device_class = DEVICE_CLASS_OUTLET
+
+    @property
+    def unique_id(self):
+        """Return the unique id of the slot."""
+        return f"{self._device.unique_id}-s{self._slot}"
+
+    @property
+    def name(self):
+        """Return the name of the switch."""
+        return f"{self._device.name} S{self._slot}"
+
+    @property
+    def assumed_state(self):
+        """Return True if unable to access real state of the switch."""
+        return False
+
+    @callback
+    def update_data(self):
+        """Update data."""
+        if self._coordinator.last_update_success:
+            self._state = self._coordinator.data[f"pwr{self._slot}"]
+        self.async_write_ha_state()
+
+    async def _async_send_packet(self, packet):
+        """Send a packet to the device."""
+        set_state = partial(self._device.api.set_state, **{f"pwr{self._slot}": packet})
+        try:
+            await self._device.async_request(set_state)
+        except (BroadlinkException, OSError) as err:
+            _LOGGER.error("Failed to send packet: %s", err)
+            return False
+        return True
diff --git a/homeassistant/components/broadlink/updater.py b/homeassistant/components/broadlink/updater.py
index eb42e688a59..c9b273218b5 100644
--- a/homeassistant/components/broadlink/updater.py
+++ b/homeassistant/components/broadlink/updater.py
@@ -26,6 +26,7 @@ def get_update_manager(device):
 
     update_managers = {
         "A1": BroadlinkA1UpdateManager,
+        "BG1": BroadlinkBG1UpdateManager,
         "MP1": BroadlinkMP1UpdateManager,
         "RM2": BroadlinkRMUpdateManager,
         "RM4": BroadlinkRMUpdateManager,
@@ -161,6 +162,14 @@ class BroadlinkSP2UpdateManager(BroadlinkUpdateManager):
         return data
 
 
+class BroadlinkBG1UpdateManager(BroadlinkUpdateManager):
+    """Manages updates for Broadlink BG1 devices."""
+
+    async def async_fetch_data(self):
+        """Fetch data from the device."""
+        return await self.device.async_request(self.device.api.get_state)
+
+
 class BroadlinkSP4UpdateManager(BroadlinkUpdateManager):
     """Manages updates for Broadlink SP4 devices."""