Add select platform to drop_connect integration (#106309)

* Add select platform to drop_connect integration

* Fix select test

* Fix minor issues

* Make function definition more specific

* Revert change to switch.py for inclusion in separate PR

* Improve typing

* Add translation keys for select options

* Fix set function typing

* Remove redundant value check

* Remove args that match defaults

homeassistant/components/drop_connect/__init__.py

@@ -15,7 +15,12 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-PLATFORMS: list[Platform] = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SWITCH]
+PLATFORMS: list[Platform] = [
+    Platform.BINARY_SENSOR,
+    Platform.SELECT,
+    Platform.SENSOR,
+    Platform.SWITCH,
+]
 
 
 async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:

homeassistant/components/drop_connect/coordinator.py

@@ -25,7 +25,7 @@ def __init__(self, hass: HomeAssistant, unique_id: str) -> None:
         super().__init__(hass, _LOGGER, name=f"{DOMAIN}-{unique_id}")
         self.drop_api = DropAPI()
 
-    async def set_water(self, value: int):
+    async def set_water(self, value: int) -> None:
         """Change water supply state."""
         payload = self.drop_api.set_water_message(value)
         await mqtt.async_publish(
@@ -34,11 +34,20 @@ async def set_water(self, value: int):
             payload,
         )
 
-    async def set_bypass(self, value: int):
+    async def set_bypass(self, value: int) -> None:
         """Change water bypass state."""
         payload = self.drop_api.set_bypass_message(value)
         await mqtt.async_publish(
             self.hass,
             self.config_entry.data[CONF_COMMAND_TOPIC],
             payload,
         )
+
+    async def set_protect_mode(self, value: str) -> None:
+        """Change protect mode state."""
+        payload = self.drop_api.set_protect_mode_message(value)
+        await mqtt.async_publish(
+            self.hass,
+            self.config_entry.data[CONF_COMMAND_TOPIC],
+            payload,
+        )

homeassistant/components/drop_connect/select.py

@@ -0,0 +1,95 @@
+"""Support for DROP selects."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+import logging
+from typing import Any
+
+from homeassistant.components.select import SelectEntity, SelectEntityDescription
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .const import CONF_DEVICE_TYPE, DEV_HUB, DOMAIN
+from .coordinator import DROPDeviceDataUpdateCoordinator
+from .entity import DROPEntity
+
+_LOGGER = logging.getLogger(__name__)
+
+# Select type constants
+PROTECT_MODE = "protect_mode"
+
+PROTECT_MODE_OPTIONS = ["away", "home", "schedule"]
+
+FLOOD_ICON = "mdi:home-flood"
+
+
+@dataclass(kw_only=True, frozen=True)
+class DROPSelectEntityDescription(SelectEntityDescription):
+    """Describes DROP select entity."""
+
+    value_fn: Callable[[DROPDeviceDataUpdateCoordinator], str | None]
+    set_fn: Callable[[DROPDeviceDataUpdateCoordinator, str], Awaitable[Any]]
+
+
+SELECTS: list[DROPSelectEntityDescription] = [
+    DROPSelectEntityDescription(
+        key=PROTECT_MODE,
+        translation_key=PROTECT_MODE,
+        icon=FLOOD_ICON,
+        options=PROTECT_MODE_OPTIONS,
+        value_fn=lambda device: device.drop_api.protect_mode(),
+        set_fn=lambda device, value: device.set_protect_mode(value),
+    )
+]
+
+# Defines which selects are used by each device type
+DEVICE_SELECTS: dict[str, list[str]] = {
+    DEV_HUB: [PROTECT_MODE],
+}
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up the DROP selects from config entry."""
+    _LOGGER.debug(
+        "Set up select for device type %s with entry_id is %s",
+        config_entry.data[CONF_DEVICE_TYPE],
+        config_entry.entry_id,
+    )
+
+    if config_entry.data[CONF_DEVICE_TYPE] in DEVICE_SELECTS:
+        async_add_entities(
+            DROPSelect(hass.data[DOMAIN][config_entry.entry_id], select)
+            for select in SELECTS
+            if select.key in DEVICE_SELECTS[config_entry.data[CONF_DEVICE_TYPE]]
+        )
+
+
+class DROPSelect(DROPEntity, SelectEntity):
+    """Representation of a DROP select."""
+
+    entity_description: DROPSelectEntityDescription
+
+    def __init__(
+        self,
+        coordinator: DROPDeviceDataUpdateCoordinator,
+        entity_description: DROPSelectEntityDescription,
+    ) -> None:
+        """Initialize the select."""
+        super().__init__(entity_description.key, coordinator)
+        self.entity_description = entity_description
+
+    @property
+    def current_option(self) -> str | None:
+        """Return the current selected option."""
+        return self.entity_description.value_fn(self.coordinator)
+
+    async def async_select_option(self, option: str) -> None:
+        """Update the current selected option."""
+        await self.entity_description.set_fn(self.coordinator, option)

homeassistant/components/drop_connect/strings.json

@@ -33,6 +33,16 @@
       "salt": { "name": "Salt low" },
       "pump": { "name": "Pump status" }
     },
+    "select": {
+      "protect_mode": {
+        "name": "Protect mode",
+        "state": {
+          "away": "Away",
+          "home": "Home",
+          "schedule": "Schedule"
+        }
+      }
+    },
     "switch": {
       "water": { "name": "Water supply" },
       "bypass": { "name": "Treatment bypass" }

tests/components/drop_connect/common.py

@@ -3,11 +3,11 @@
 TEST_DATA_HUB_TOPIC = "drop_connect/DROP-1_C0FFEE/255"
 TEST_DATA_HUB = (
     '{"curFlow":5.77,"peakFlow":13.8,"usedToday":232.77,"avgUsed":76,"psi":62.2,"psiLow":61,"psiHigh":62,'
-    '"water":1,"bypass":0,"pMode":"HOME","battery":50,"notif":1,"leak":0}'
+    '"water":1,"bypass":0,"pMode":"home","battery":50,"notif":1,"leak":0}'
 )
 TEST_DATA_HUB_RESET = (
     '{"curFlow":0,"peakFlow":0,"usedToday":0,"avgUsed":0,"psi":0,"psiLow":0,"psiHigh":0,'
-    '"water":0,"bypass":1,"pMode":"AWAY","battery":0,"notif":0,"leak":0}'
+    '"water":0,"bypass":1,"pMode":"away","battery":0,"notif":0,"leak":0}'
 )
 
 TEST_DATA_SALT_TOPIC = "drop_connect/DROP-1_C0FFEE/8"

tests/components/drop_connect/test_select.py

@@ -0,0 +1,59 @@
+"""Test DROP select entities."""
+
+from homeassistant.components.drop_connect.const import DOMAIN
+from homeassistant.components.select import (
+    ATTR_OPTION,
+    ATTR_OPTIONS,
+    DOMAIN as SELECT_DOMAIN,
+    SERVICE_SELECT_OPTION,
+)
+from homeassistant.const import ATTR_ENTITY_ID
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from .common import TEST_DATA_HUB, TEST_DATA_HUB_RESET, TEST_DATA_HUB_TOPIC
+
+from tests.common import async_fire_mqtt_message
+from tests.typing import MqttMockHAClient
+
+
+async def test_selects_hub(
+    hass: HomeAssistant, config_entry_hub, mqtt_mock: MqttMockHAClient
+) -> None:
+    """Test DROP binary sensors for hubs."""
+    config_entry_hub.add_to_hass(hass)
+    assert await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    protect_mode_select_name = "select.hub_drop_1_c0ffee_protect_mode"
+    protect_mode_select = hass.states.get(protect_mode_select_name)
+    assert protect_mode_select
+    assert protect_mode_select.attributes.get(ATTR_OPTIONS) == [
+        "away",
+        "home",
+        "schedule",
+    ]
+
+    async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB)
+    await hass.async_block_till_done()
+
+    protect_mode_select = hass.states.get(protect_mode_select_name)
+    assert protect_mode_select
+    assert protect_mode_select.state == "home"
+
+    await hass.services.async_call(
+        SELECT_DOMAIN,
+        SERVICE_SELECT_OPTION,
+        {ATTR_OPTION: "away", ATTR_ENTITY_ID: protect_mode_select_name},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+
+    async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
+    await hass.async_block_till_done()
+
+    protect_mode_select = hass.states.get(protect_mode_select_name)
+    assert protect_mode_select
+    assert protect_mode_select.state == "away"