Add config flow to Switchbot (#50653)

Co-authored-by: Daniel Hjelseth Høyer <mail@dahoiv.net>
Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/switchbot/switch.py

@@ -1,18 +1,48 @@
-"""Support for Switchbot."""
+"""Support for Switchbot bot."""
 from __future__ import annotations
 
+import logging
 from typing import Any
 
-# pylint: disable=import-error
-import switchbot
 import voluptuous as vol
 
-from homeassistant.components.switch import PLATFORM_SCHEMA, SwitchEntity
-from homeassistant.const import CONF_MAC, CONF_NAME, CONF_PASSWORD
-import homeassistant.helpers.config_validation as cv
+from homeassistant.components.switch import (
+    DEVICE_CLASS_SWITCH,
+    PLATFORM_SCHEMA,
+    SwitchEntity,
+)
+from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
+from homeassistant.const import (
+    CONF_MAC,
+    CONF_NAME,
+    CONF_PASSWORD,
+    CONF_SENSOR_TYPE,
+    STATE_ON,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import (
+    config_validation as cv,
+    device_registry as dr,
+    entity_platform,
+)
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.restore_state import RestoreEntity
+from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-DEFAULT_NAME = "Switchbot"
+from .const import (
+    ATTR_BOT,
+    CONF_RETRY_COUNT,
+    DATA_COORDINATOR,
+    DEFAULT_NAME,
+    DOMAIN,
+    MANUFACTURER,
+)
+from .coordinator import SwitchbotDataUpdateCoordinator
+
+# Initialize the logger
+_LOGGER = logging.getLogger(__name__)
+PARALLEL_UPDATES = 1
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
@@ -23,46 +53,120 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
 )
 
 
-def setup_platform(hass, config, add_entities, discovery_info=None):
-    """Perform the setup for Switchbot devices."""
-    name = config.get(CONF_NAME)
-    mac_addr = config[CONF_MAC]
-    password = config.get(CONF_PASSWORD)
-    add_entities([SwitchBot(mac_addr, name, password)])
+async def async_setup_platform(
+    hass: HomeAssistant,
+    config: ConfigType,
+    async_add_entities: entity_platform.AddEntitiesCallback,
+    discovery_info: DiscoveryInfoType | None = None,
+) -> None:
+    """Import yaml config and initiates config flow for Switchbot devices."""
+
+    # Check if entry config exists and skips import if it does.
+    if hass.config_entries.async_entries(DOMAIN):
+        return
+
+    hass.async_create_task(
+        hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": SOURCE_IMPORT},
+            data={
+                CONF_NAME: config[CONF_NAME],
+                CONF_PASSWORD: config.get(CONF_PASSWORD, None),
+                CONF_MAC: config[CONF_MAC].replace("-", ":").lower(),
+                CONF_SENSOR_TYPE: ATTR_BOT,
+            },
+        )
+    )
 
 
-class SwitchBot(SwitchEntity, RestoreEntity):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: entity_platform.AddEntitiesCallback,
+) -> None:
+    """Set up Switchbot based on a config entry."""
+    coordinator: SwitchbotDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][
+        DATA_COORDINATOR
+    ]
+
+    if entry.data[CONF_SENSOR_TYPE] != ATTR_BOT:
+        return
+
+    async_add_entities(
+        [
+            SwitchBot(
+                coordinator,
+                entry.unique_id,
+                entry.data[CONF_MAC],
+                entry.data[CONF_NAME],
+                entry.data.get(CONF_PASSWORD, None),
+                entry.options[CONF_RETRY_COUNT],
+            )
+        ]
+    )
+
+
+class SwitchBot(CoordinatorEntity, SwitchEntity, RestoreEntity):
     """Representation of a Switchbot."""
 
-    def __init__(self, mac, name, password) -> None:
+    coordinator: SwitchbotDataUpdateCoordinator
+    _attr_device_class = DEVICE_CLASS_SWITCH
+
+    def __init__(
+        self,
+        coordinator: SwitchbotDataUpdateCoordinator,
+        idx: str | None,
+        mac: str,
+        name: str,
+        password: str,
+        retry_count: int,
+    ) -> None:
         """Initialize the Switchbot."""
-
-        self._state: bool | None = None
+        super().__init__(coordinator)
+        self._idx = idx
         self._last_run_success: bool | None = None
-        self._name = name
         self._mac = mac
-        self._device = switchbot.Switchbot(mac=mac, password=password)
+        self._device = self.coordinator.switchbot_api.Switchbot(
+            mac=mac, password=password, retry_count=retry_count
+        )
+        self._attr_unique_id = self._mac.replace(":", "")
+        self._attr_name = name
+        self._attr_device_info: DeviceInfo = {
+            "connections": {(dr.CONNECTION_NETWORK_MAC, self._mac)},
+            "name": name,
+            "model": self.coordinator.data[self._idx]["modelName"],
+            "manufacturer": MANUFACTURER,
+        }
 
-    async def async_added_to_hass(self):
+    async def async_added_to_hass(self) -> None:
         """Run when entity about to be added."""
         await super().async_added_to_hass()
-        state = await self.async_get_last_state()
-        if not state:
+        last_state = await self.async_get_last_state()
+        if not last_state:
             return
-        self._state = state.state == "on"
+        self._attr_is_on = last_state.state == STATE_ON
+        self._last_run_success = last_state.attributes["last_run_success"]
 
-    def turn_on(self, **kwargs) -> None:
+    async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn device on."""
-        if self._device.turn_on():
-            self._state = True
+        _LOGGER.info("Turn Switchbot bot on %s", self._mac)
+
+        async with self.coordinator.api_lock:
+            update_ok = await self.hass.async_add_executor_job(self._device.turn_on)
+
+        if update_ok:
             self._last_run_success = True
         else:
             self._last_run_success = False
 
-    def turn_off(self, **kwargs) -> None:
+    async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn device off."""
-        if self._device.turn_off():
-            self._state = False
+        _LOGGER.info("Turn Switchbot bot off %s", self._mac)
+
+        async with self.coordinator.api_lock:
+            update_ok = await self.hass.async_add_executor_job(self._device.turn_off)
+
+        if update_ok:
             self._last_run_success = True
         else:
             self._last_run_success = False
@@ -70,24 +174,20 @@ class SwitchBot(SwitchEntity, RestoreEntity):
     @property
     def assumed_state(self) -> bool:
         """Return true if unable to access real state of entity."""
-        return True
+        if not self.coordinator.data[self._idx]["data"]["switchMode"]:
+            return True
+        return False
 
     @property
     def is_on(self) -> bool:
         """Return true if device is on."""
-        return bool(self._state)
+        return self.coordinator.data[self._idx]["data"]["isOn"]
 
     @property
-    def unique_id(self) -> str:
-        """Return a unique, Home Assistant friendly identifier for this entity."""
-        return self._mac.replace(":", "")
-
-    @property
-    def name(self) -> str:
-        """Return the name of the switch."""
-        return self._name
-
-    @property
-    def extra_state_attributes(self) -> dict[str, Any]:
+    def device_state_attributes(self) -> dict:
         """Return the state attributes."""
-        return {"last_run_success": self._last_run_success}
+        return {
+            "last_run_success": self._last_run_success,
+            "mac_address": self._mac,
+            "switch_mode": self.coordinator.data[self._idx]["data"]["switchMode"],
+        }