Add cover platform to switchbot (#56414)

Co-authored-by: J. Nick Koston <nick@koston.org>

.coveragerc

@@ -1014,6 +1014,8 @@ omit =
     homeassistant/components/switchbot/switch.py
     homeassistant/components/switchbot/__init__.py
     homeassistant/components/switchbot/const.py
+    homeassistant/components/switchbot/entity.py
+    homeassistant/components/switchbot/cover.py
     homeassistant/components/switchbot/coordinator.py
     homeassistant/components/switchmate/switch.py
     homeassistant/components/syncthing/__init__.py

homeassistant/components/switchbot/__init__.py

@@ -4,10 +4,13 @@ from asyncio import Lock
 import switchbot  # pylint: disable=import-error
 
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_SENSOR_TYPE
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 
 from .const import (
+    ATTR_BOT,
+    ATTR_CURTAIN,
     BTLE_LOCK,
     COMMON_OPTIONS,
     CONF_RETRY_COUNT,
@@ -23,7 +26,10 @@ from .const import (
 )
 from .coordinator import SwitchbotDataUpdateCoordinator
 
-PLATFORMS = ["switch"]
+PLATFORMS_BY_TYPE = {
+    ATTR_BOT: ["switch"],
+    ATTR_CURTAIN: ["cover"],
+}
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -83,14 +89,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     hass.data[DOMAIN][entry.entry_id] = {DATA_COORDINATOR: coordinator}
 
-    hass.config_entries.async_setup_platforms(entry, PLATFORMS)
+    sensor_type = entry.data[CONF_SENSOR_TYPE]
+
+    hass.config_entries.async_setup_platforms(entry, PLATFORMS_BY_TYPE[sensor_type])
 
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    sensor_type = entry.data[CONF_SENSOR_TYPE]
+    unload_ok = await hass.config_entries.async_unload_platforms(
+        entry, PLATFORMS_BY_TYPE[sensor_type]
+    )
 
     if unload_ok:
         hass.data[DOMAIN].pop(entry.entry_id)

homeassistant/components/switchbot/config_flow.py

@@ -14,7 +14,6 @@ from homeassistant.core import callback
 from homeassistant.data_entry_flow import FlowResult
 
 from .const import (
-    ATTR_BOT,
     BTLE_LOCK,
     CONF_RETRY_COUNT,
     CONF_RETRY_TIMEOUT,
@@ -25,6 +24,7 @@ from .const import (
     DEFAULT_SCAN_TIMEOUT,
     DEFAULT_TIME_BETWEEN_UPDATE_COMMAND,
     DOMAIN,
+    SUPPORTED_MODEL_TYPES,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -70,8 +70,8 @@ class SwitchbotConfigFlow(ConfigFlow, domain=DOMAIN):
                 _btle_connect, data[CONF_MAC]
             )
 
-        if _btle_adv_data["modelName"] == "WoHand":
+        if _btle_adv_data["modelName"] in SUPPORTED_MODEL_TYPES:
-            data[CONF_SENSOR_TYPE] = ATTR_BOT
+            data[CONF_SENSOR_TYPE] = SUPPORTED_MODEL_TYPES[_btle_adv_data["modelName"]]
             return self.async_create_entry(title=data[CONF_NAME], data=data)
 
         return self.async_abort(reason="switchbot_unsupported_type")

homeassistant/components/switchbot/const.py

@@ -4,7 +4,9 @@ MANUFACTURER = "switchbot"
 
 # Config Attributes
 ATTR_BOT = "bot"
+ATTR_CURTAIN = "curtain"
 DEFAULT_NAME = "Switchbot"
+SUPPORTED_MODEL_TYPES = {"WoHand": ATTR_BOT, "WoCurtain": ATTR_CURTAIN}
 
 # Config Defaults
 DEFAULT_RETRY_COUNT = 3

homeassistant/components/switchbot/cover.py

@@ -0,0 +1,141 @@
+"""Support for SwitchBot curtains."""
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from switchbot import SwitchbotCurtain  # pylint: disable=import-error
+
+from homeassistant.components.cover import (
+    ATTR_CURRENT_POSITION,
+    ATTR_POSITION,
+    DEVICE_CLASS_CURTAIN,
+    SUPPORT_CLOSE,
+    SUPPORT_OPEN,
+    SUPPORT_SET_POSITION,
+    SUPPORT_STOP,
+    CoverEntity,
+)
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_MAC, CONF_NAME, CONF_PASSWORD
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
+
+from .const import CONF_RETRY_COUNT, DATA_COORDINATOR, DOMAIN
+from .coordinator import SwitchbotDataUpdateCoordinator
+from .entity import SwitchbotEntity
+
+# Initialize the logger
+_LOGGER = logging.getLogger(__name__)
+PARALLEL_UPDATES = 1
+
+
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
+    """Set up Switchbot curtain based on a config entry."""
+    coordinator: SwitchbotDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][
+        DATA_COORDINATOR
+    ]
+
+    async_add_entities(
+        [
+            SwitchBotCurtainEntity(
+                coordinator,
+                entry.unique_id,
+                entry.data[CONF_MAC],
+                entry.data[CONF_NAME],
+                coordinator.switchbot_api.SwitchbotCurtain(
+                    mac=entry.data[CONF_MAC],
+                    password=entry.data.get(CONF_PASSWORD),
+                    retry_count=entry.options[CONF_RETRY_COUNT],
+                ),
+            )
+        ]
+    )
+
+
+class SwitchBotCurtainEntity(SwitchbotEntity, CoverEntity, RestoreEntity):
+    """Representation of a Switchbot."""
+
+    coordinator: SwitchbotDataUpdateCoordinator
+    _attr_device_class = DEVICE_CLASS_CURTAIN
+    _attr_supported_features = (
+        SUPPORT_OPEN | SUPPORT_CLOSE | SUPPORT_STOP | SUPPORT_SET_POSITION
+    )
+    _attr_assumed_state = True
+
+    def __init__(
+        self,
+        coordinator: SwitchbotDataUpdateCoordinator,
+        idx: str | None,
+        mac: str,
+        name: str,
+        device: SwitchbotCurtain,
+    ) -> None:
+        """Initialize the Switchbot."""
+        super().__init__(coordinator, idx, mac, name)
+        self._attr_unique_id = idx
+        self._device = device
+
+    async def async_added_to_hass(self) -> None:
+        """Run when entity about to be added."""
+        await super().async_added_to_hass()
+        last_state = await self.async_get_last_state()
+        if not last_state or ATTR_CURRENT_POSITION not in last_state.attributes:
+            return
+
+        self._attr_current_cover_position = last_state.attributes[ATTR_CURRENT_POSITION]
+        self._last_run_success = last_state.attributes["last_run_success"]
+        self._attr_is_closed = last_state.attributes[ATTR_CURRENT_POSITION] <= 20
+
+    async def async_open_cover(self, **kwargs: Any) -> None:
+        """Open the curtain."""
+
+        _LOGGER.debug("Switchbot to open curtain %s", self._mac)
+
+        async with self.coordinator.api_lock:
+            self._last_run_success = bool(
+                await self.hass.async_add_executor_job(self._device.open)
+            )
+
+    async def async_close_cover(self, **kwargs: Any) -> None:
+        """Close the curtain."""
+
+        _LOGGER.debug("Switchbot to close the curtain %s", self._mac)
+
+        async with self.coordinator.api_lock:
+            self._last_run_success = bool(
+                await self.hass.async_add_executor_job(self._device.close)
+            )
+
+    async def async_stop_cover(self, **kwargs: Any) -> None:
+        """Stop the moving of this device."""
+
+        _LOGGER.debug("Switchbot to stop %s", self._mac)
+
+        async with self.coordinator.api_lock:
+            self._last_run_success = bool(
+                await self.hass.async_add_executor_job(self._device.stop)
+            )
+
+    async def async_set_cover_position(self, **kwargs: Any) -> None:
+        """Move the cover shutter to a specific position."""
+        position = kwargs.get(ATTR_POSITION)
+
+        _LOGGER.debug("Switchbot to move at %d %s", position, self._mac)
+
+        async with self.coordinator.api_lock:
+            self._last_run_success = bool(
+                await self.hass.async_add_executor_job(
+                    self._device.set_position, position
+                )
+            )
+
+    @callback
+    def _handle_coordinator_update(self) -> None:
+        """Handle updated data from the coordinator."""
+        self._attr_current_cover_position = self.data["data"]["position"]
+        self._attr_is_closed = self.data["data"]["position"] <= 20
+        self.async_write_ha_state()

homeassistant/components/switchbot/entity.py

@@ -0,0 +1,46 @@
+"""An abstract class common to all Switchbot entities."""
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers.entity import DeviceInfo, Entity
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .const import MANUFACTURER
+from .coordinator import SwitchbotDataUpdateCoordinator
+
+
+class SwitchbotEntity(CoordinatorEntity, Entity):
+    """Generic entity encapsulating common features of Switchbot device."""
+
+    def __init__(
+        self,
+        coordinator: SwitchbotDataUpdateCoordinator,
+        idx: str | None,
+        mac: str,
+        name: str,
+    ) -> None:
+        """Initialize the entity."""
+        super().__init__(coordinator)
+        self._last_run_success: bool | None = None
+        self._idx = idx
+        self._mac = mac
+        self._attr_name = name
+        self._attr_device_info: DeviceInfo = {
+            "connections": {(dr.CONNECTION_NETWORK_MAC, self._mac)},
+            "name": self._attr_name,
+            "model": self.data["modelName"],
+            "manufacturer": MANUFACTURER,
+        }
+
+    @property
+    def data(self) -> dict[str, Any]:
+        """Return coordinator data for this entity."""
+        return self.coordinator.data[self._idx]
+
+    @property
+    def extra_state_attributes(self) -> Mapping[Any, Any]:
+        """Return the state attributes."""
+        return {"last_run_success": self._last_run_success, "mac_address": self._mac}

homeassistant/components/switchbot/switch.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 import logging
 from typing import Any
 
+from switchbot import Switchbot  # pylint: disable=import-error
 import voluptuous as vol
 
 from homeassistant.components.switch import (
@@ -20,25 +21,13 @@ from homeassistant.const import (
     STATE_ON,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import (
+from homeassistant.helpers import config_validation as cv, entity_platform
-    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
 
-from .const import (
+from .const import ATTR_BOT, CONF_RETRY_COUNT, DATA_COORDINATOR, DEFAULT_NAME, DOMAIN
-    ATTR_BOT,
-    CONF_RETRY_COUNT,
-    DATA_COORDINATOR,
-    DEFAULT_NAME,
-    DOMAIN,
-    MANUFACTURER,
-)
 from .coordinator import SwitchbotDataUpdateCoordinator
+from .entity import SwitchbotEntity
 
 # Initialize the logger
 _LOGGER = logging.getLogger(__name__)
@@ -89,24 +78,24 @@ async def async_setup_entry(
         DATA_COORDINATOR
     ]
 
-    if entry.data[CONF_SENSOR_TYPE] != ATTR_BOT:
-        return
-
     async_add_entities(
         [
-            SwitchBot(
+            SwitchBotBotEntity(
                 coordinator,
                 entry.unique_id,
                 entry.data[CONF_MAC],
                 entry.data[CONF_NAME],
-                entry.data.get(CONF_PASSWORD, None),
+                coordinator.switchbot_api.Switchbot(
-                entry.options[CONF_RETRY_COUNT],
+                    mac=entry.data[CONF_MAC],
+                    password=entry.data.get(CONF_PASSWORD),
+                    retry_count=entry.options[CONF_RETRY_COUNT],
+                ),
             )
         ]
     )
 
 
-class SwitchBot(CoordinatorEntity, SwitchEntity, RestoreEntity):
+class SwitchBotBotEntity(SwitchbotEntity, SwitchEntity, RestoreEntity):
     """Representation of a Switchbot."""
 
     coordinator: SwitchbotDataUpdateCoordinator
@@ -118,25 +107,12 @@ class SwitchBot(CoordinatorEntity, SwitchEntity, RestoreEntity):
         idx: str | None,
         mac: str,
         name: str,
-        password: str,
+        device: Switchbot,
-        retry_count: int,
     ) -> None:
         """Initialize the Switchbot."""
-        super().__init__(coordinator)
+        super().__init__(coordinator, idx, mac, name)
-        self._idx = idx
-        self._last_run_success: bool | None = None
-        self._mac = mac
-        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._device = device
-        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) -> None:
         """Run when entity about to be added."""
@@ -152,42 +128,35 @@ class SwitchBot(CoordinatorEntity, SwitchEntity, RestoreEntity):
         _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)
+            self._last_run_success = bool(
-
+                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
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn device off."""
         _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)
+            self._last_run_success = bool(
-
+                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
 
     @property
     def assumed_state(self) -> bool:
         """Return true if unable to access real state of entity."""
-        if not self.coordinator.data[self._idx]["data"]["switchMode"]:
+        if not self.data["data"]["switchMode"]:
             return True
         return False
 
     @property
     def is_on(self) -> bool:
         """Return true if device is on."""
-        return self.coordinator.data[self._idx]["data"]["isOn"]
+        return self.data["data"]["isOn"]
 
     @property
-    def device_state_attributes(self) -> dict:
+    def extra_state_attributes(self) -> dict:
         """Return the state attributes."""
         return {
-            "last_run_success": self._last_run_success,
+            **super().extra_state_attributes,
-            "mac_address": self._mac,
+            "switch_mode": self.data["data"]["switchMode"],
-            "switch_mode": self.coordinator.data[self._idx]["data"]["switchMode"],
         }

tests/components/switchbot/__init__.py

@@ -20,6 +20,12 @@ USER_INPUT = {
     CONF_MAC: "e7:89:43:99:99:99",
 }
 
+USER_INPUT_CURTAIN = {
+    CONF_NAME: "test-name",
+    CONF_PASSWORD: "test-password",
+    CONF_MAC: "e7:89:43:90:90:90",
+}
+
 USER_INPUT_UNSUPPORTED_DEVICE = {
     CONF_NAME: "test-name",
     CONF_PASSWORD: "test-password",

tests/components/switchbot/conftest.py

@@ -25,6 +25,21 @@ class MocGetSwitchbotDevices:
             "model": "H",
             "modelName": "WoHand",
         }
+        self._curtain_all_services_data = {
+            "mac_address": "e7:89:43:90:90:90",
+            "Flags": "06",
+            "Manufacturer": "5900e78943d9fe7c",
+            "Complete 128b Services": "cba20d00-224d-11e6-9fb8-0002a5d5c51b",
+            "data": {
+                "calibration": True,
+                "battery": 74,
+                "position": 100,
+                "lightLevel": 2,
+                "rssi": -73,
+            },
+            "model": "c",
+            "modelName": "WoCurtain",
+        }
         self._unsupported_device = {
             "mac_address": "test",
             "Flags": "06",
@@ -50,6 +65,8 @@ class MocGetSwitchbotDevices:
             return self._all_services_data
         if mac == "test":
             return self._unsupported_device
+        if mac == "e7:89:43:90:90:90":
+            return self._curtain_all_services_data
 
         return None
 

tests/components/switchbot/test_config_flow.py

@@ -1,8 +1,12 @@
 """Test the switchbot config flow."""
 
-from unittest.mock import patch
-
 from homeassistant.components.switchbot.config_flow import NotConnectedError
+from homeassistant.components.switchbot.const import (
+    CONF_RETRY_COUNT,
+    CONF_RETRY_TIMEOUT,
+    CONF_SCAN_TIMEOUT,
+    CONF_TIME_BETWEEN_UPDATE_COMMAND,
+)
 from homeassistant.config_entries import SOURCE_IMPORT, SOURCE_USER
 from homeassistant.const import CONF_MAC, CONF_NAME, CONF_PASSWORD, CONF_SENSOR_TYPE
 from homeassistant.data_entry_flow import (
@@ -14,6 +18,7 @@ from homeassistant.setup import async_setup_component
 
 from . import (
     USER_INPUT,
+    USER_INPUT_CURTAIN,
     USER_INPUT_INVALID,
     USER_INPUT_UNSUPPORTED_DEVICE,
     YAML_CONFIG,
@@ -71,6 +76,33 @@ async def test_user_form_valid_mac(hass):
     assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "already_configured_device"
 
+    # test curtain device creation.
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": SOURCE_USER}
+    )
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == "user"
+    assert result["errors"] == {}
+
+    with _patch_async_setup_entry() as mock_setup_entry:
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            USER_INPUT_CURTAIN,
+        )
+    await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_CREATE_ENTRY
+    assert result["title"] == "test-name"
+    assert result["data"] == {
+        CONF_MAC: "e7:89:43:90:90:90",
+        CONF_NAME: "test-name",
+        CONF_PASSWORD: "test-password",
+        CONF_SENSOR_TYPE: "curtain",
+    }
+
+    assert len(mock_setup_entry.mock_calls) == 1
+
 
 async def test_user_form_unsupported_device(hass):
     """Test the user initiated form for unsupported device type."""
@@ -165,62 +197,58 @@ async def test_user_form_exception(hass, switchbot_config_flow):
 
 async def test_options_flow(hass):
     """Test updating options."""
-    with patch("homeassistant.components.switchbot.PLATFORMS", []):
+    with _patch_async_setup_entry() as mock_setup_entry:
         entry = await init_integration(hass)
 
-    assert entry.options["update_time"] == 60
+        result = await hass.config_entries.options.async_init(entry.entry_id)
-    assert entry.options["retry_count"] == 3
+        assert result["type"] == RESULT_TYPE_FORM
-    assert entry.options["retry_timeout"] == 5
+        assert result["step_id"] == "init"
-    assert entry.options["scan_timeout"] == 5
+        assert result["errors"] is None
 
-    result = await hass.config_entries.options.async_init(entry.entry_id)
-    assert result["type"] == RESULT_TYPE_FORM
-    assert result["step_id"] == "init"
-    assert result["errors"] is None
-
-    with _patch_async_setup_entry() as mock_setup_entry:
         result = await hass.config_entries.options.async_configure(
             result["flow_id"],
             user_input={
-                "update_time": 60,
+                CONF_TIME_BETWEEN_UPDATE_COMMAND: 60,
-                "retry_count": 3,
+                CONF_RETRY_COUNT: 3,
-                "retry_timeout": 5,
+                CONF_RETRY_TIMEOUT: 5,
-                "scan_timeout": 5,
+                CONF_SCAN_TIMEOUT: 5,
             },
         )
-    await hass.async_block_till_done()
+        await hass.async_block_till_done()
 
     assert result["type"] == RESULT_TYPE_CREATE_ENTRY
-    assert result["data"]["update_time"] == 60
+    assert result["data"][CONF_TIME_BETWEEN_UPDATE_COMMAND] == 60
-    assert result["data"]["retry_count"] == 3
+    assert result["data"][CONF_RETRY_COUNT] == 3
-    assert result["data"]["retry_timeout"] == 5
+    assert result["data"][CONF_RETRY_TIMEOUT] == 5
-    assert result["data"]["scan_timeout"] == 5
+    assert result["data"][CONF_SCAN_TIMEOUT] == 5
 
-    assert len(mock_setup_entry.mock_calls) == 0
+    assert len(mock_setup_entry.mock_calls) == 1
 
     # Test changing of entry options.
 
-    result = await hass.config_entries.options.async_init(entry.entry_id)
-    assert result["type"] == RESULT_TYPE_FORM
-    assert result["step_id"] == "init"
-    assert result["errors"] is None
-
     with _patch_async_setup_entry() as mock_setup_entry:
+        entry = await init_integration(hass)
+
+        result = await hass.config_entries.options.async_init(entry.entry_id)
+        assert result["type"] == RESULT_TYPE_FORM
+        assert result["step_id"] == "init"
+        assert result["errors"] is None
+
         result = await hass.config_entries.options.async_configure(
             result["flow_id"],
             user_input={
-                "update_time": 60,
+                CONF_TIME_BETWEEN_UPDATE_COMMAND: 66,
-                "retry_count": 3,
+                CONF_RETRY_COUNT: 6,
-                "retry_timeout": 5,
+                CONF_RETRY_TIMEOUT: 6,
-                "scan_timeout": 5,
+                CONF_SCAN_TIMEOUT: 6,
             },
         )
-    await hass.async_block_till_done()
+        await hass.async_block_till_done()
 
     assert result["type"] == RESULT_TYPE_CREATE_ENTRY
-    assert result["data"]["update_time"] == 60
+    assert result["data"][CONF_TIME_BETWEEN_UPDATE_COMMAND] == 66
-    assert result["data"]["retry_count"] == 3
+    assert result["data"][CONF_RETRY_COUNT] == 6
-    assert result["data"]["retry_timeout"] == 5
+    assert result["data"][CONF_RETRY_TIMEOUT] == 6
-    assert result["data"]["scan_timeout"] == 5
+    assert result["data"][CONF_SCAN_TIMEOUT] == 6
 
-    assert len(mock_setup_entry.mock_calls) == 0
+    assert len(mock_setup_entry.mock_calls) == 1