Add switches to drop_connect integration (#106264)

* Add switches to drop_connect integration

* Let device update state

* Simplify icon property

* Use constants for icon names

* Add simulated responses from devices

* Use keyed form for switch value

* Clean up properties

homeassistant/components/drop_connect/__init__.py

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

homeassistant/components/drop_connect/coordinator.py

@@ -5,11 +5,12 @@ import logging
 
 from dropmqttapi.mqttapi import DropAPI
 
+from homeassistant.components import mqtt
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import DOMAIN
+from .const import CONF_COMMAND_TOPIC, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -23,3 +24,21 @@ class DROPDeviceDataUpdateCoordinator(DataUpdateCoordinator):
         """Initialize the device."""
         super().__init__(hass, _LOGGER, name=f"{DOMAIN}-{unique_id}")
         self.drop_api = DropAPI()
+
+    async def set_water(self, value: int):
+        """Change water supply state."""
+        payload = self.drop_api.set_water_message(value)
+        await mqtt.async_publish(
+            self.hass,
+            self.config_entry.data[CONF_COMMAND_TOPIC],
+            payload,
+        )
+
+    async def set_bypass(self, value: int):
+        """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,
+        )

homeassistant/components/drop_connect/strings.json

@@ -32,6 +32,10 @@
       "reserve_in_use": { "name": "Reserve capacity in use" },
       "salt": { "name": "Salt low" },
       "pump": { "name": "Pump status" }
+    },
+    "switch": {
+      "water": { "name": "Water supply" },
+      "bypass": { "name": "Treatment bypass" }
     }
   }
 }

homeassistant/components/drop_connect/switch.py

@@ -0,0 +1,124 @@
+"""Support for DROP switches."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+import logging
+from typing import Any
+
+from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
+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_FILTER,
+    DEV_HUB,
+    DEV_PROTECTION_VALVE,
+    DEV_SOFTENER,
+    DOMAIN,
+)
+from .coordinator import DROPDeviceDataUpdateCoordinator
+from .entity import DROPEntity
+
+_LOGGER = logging.getLogger(__name__)
+
+ICON_VALVE_OPEN = "mdi:valve-open"
+ICON_VALVE_CLOSED = "mdi:valve-closed"
+ICON_VALVE_UNKNOWN = "mdi:valve"
+ICON_VALVE = {False: ICON_VALVE_CLOSED, True: ICON_VALVE_OPEN, None: ICON_VALVE_UNKNOWN}
+
+SWITCH_VALUE: dict[int | None, bool] = {0: False, 1: True}
+
+# Switch type constants
+WATER_SWITCH = "water"
+BYPASS_SWITCH = "bypass"
+
+
+@dataclass(kw_only=True, frozen=True)
+class DROPSwitchEntityDescription(SwitchEntityDescription):
+    """Describes DROP switch entity."""
+
+    value_fn: Callable[[DROPDeviceDataUpdateCoordinator], int | None]
+    set_fn: Callable
+
+
+SWITCHES: list[DROPSwitchEntityDescription] = [
+    DROPSwitchEntityDescription(
+        key=WATER_SWITCH,
+        translation_key=WATER_SWITCH,
+        icon=ICON_VALVE_UNKNOWN,
+        value_fn=lambda device: device.drop_api.water(),
+        set_fn=lambda device, value: device.set_water(value),
+    ),
+    DROPSwitchEntityDescription(
+        key=BYPASS_SWITCH,
+        translation_key=BYPASS_SWITCH,
+        icon=ICON_VALVE_UNKNOWN,
+        value_fn=lambda device: device.drop_api.bypass(),
+        set_fn=lambda device, value: device.set_bypass(value),
+    ),
+]
+
+# Defines which switches are used by each device type
+DEVICE_SWITCHES: dict[str, list[str]] = {
+    DEV_FILTER: [BYPASS_SWITCH],
+    DEV_HUB: [WATER_SWITCH, BYPASS_SWITCH],
+    DEV_PROTECTION_VALVE: [WATER_SWITCH],
+    DEV_SOFTENER: [BYPASS_SWITCH],
+}
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up the DROP switches from config entry."""
+    _LOGGER.debug(
+        "Set up switch 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_SWITCHES:
+        async_add_entities(
+            DROPSwitch(hass.data[DOMAIN][config_entry.entry_id], switch)
+            for switch in SWITCHES
+            if switch.key in DEVICE_SWITCHES[config_entry.data[CONF_DEVICE_TYPE]]
+        )
+
+
+class DROPSwitch(DROPEntity, SwitchEntity):
+    """Representation of a DROP switch."""
+
+    entity_description: DROPSwitchEntityDescription
+
+    def __init__(
+        self,
+        coordinator: DROPDeviceDataUpdateCoordinator,
+        entity_description: DROPSwitchEntityDescription,
+    ) -> None:
+        """Initialize the switch."""
+        super().__init__(entity_description.key, coordinator)
+        self.entity_description = entity_description
+
+    @property
+    def is_on(self) -> bool | None:
+        """Return the state of the binary sensor."""
+        return SWITCH_VALUE.get(self.entity_description.value_fn(self.coordinator))
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn switch on."""
+        await self.entity_description.set_fn(self.coordinator, 1)
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn switch off."""
+        await self.entity_description.set_fn(self.coordinator, 0)
+
+    @property
+    def icon(self) -> str:
+        """Return the icon to use for dynamic states."""
+        return ICON_VALVE[self.is_on]

tests/components/drop_connect/common.py

@@ -7,7 +7,7 @@ TEST_DATA_HUB = (
 )
 TEST_DATA_HUB_RESET = (
     '{"curFlow":0,"peakFlow":0,"usedToday":0,"avgUsed":0,"psi":0,"psiLow":0,"psiHigh":0,'
-    '"water":0,"bypass":0,"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"
@@ -23,12 +23,12 @@ TEST_DATA_SOFTENER = (
     '{"curFlow":5.0,"bypass":0,"battery":20,"capacity":1000,"resInUse":1,"psi":50.5}'
 )
 TEST_DATA_SOFTENER_RESET = (
-    '{"curFlow":0,"bypass":0,"battery":0,"capacity":0,"resInUse":0,"psi":null}'
+    '{"curFlow":0,"bypass":1,"battery":0,"capacity":0,"resInUse":0,"psi":null}'
 )
 
 TEST_DATA_FILTER_TOPIC = "drop_connect/DROP-1_C0FFEE/4"
 TEST_DATA_FILTER = '{"curFlow":19.84,"bypass":0,"battery":12,"psi":38.2}'
-TEST_DATA_FILTER_RESET = '{"curFlow":0,"bypass":0,"battery":0,"psi":null}'
+TEST_DATA_FILTER_RESET = '{"curFlow":0,"bypass":1,"battery":0,"psi":null}'
 
 TEST_DATA_PROTECTION_VALVE_TOPIC = "drop_connect/DROP-1_C0FFEE/78"
 TEST_DATA_PROTECTION_VALVE = (

tests/components/drop_connect/test_switch.py

@@ -0,0 +1,275 @@
+"""Test DROP switch entities."""
+
+from homeassistant.components.drop_connect.const import DOMAIN
+from homeassistant.components.switch import (
+    DOMAIN as SWITCH_DOMAIN,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+)
+from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON, STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from .common import (
+    TEST_DATA_FILTER,
+    TEST_DATA_FILTER_RESET,
+    TEST_DATA_FILTER_TOPIC,
+    TEST_DATA_HUB,
+    TEST_DATA_HUB_RESET,
+    TEST_DATA_HUB_TOPIC,
+    TEST_DATA_PROTECTION_VALVE,
+    TEST_DATA_PROTECTION_VALVE_RESET,
+    TEST_DATA_PROTECTION_VALVE_TOPIC,
+    TEST_DATA_SOFTENER,
+    TEST_DATA_SOFTENER_RESET,
+    TEST_DATA_SOFTENER_TOPIC,
+)
+
+from tests.common import async_fire_mqtt_message
+from tests.typing import MqttMockHAClient
+
+
+async def test_switches_hub(
+    hass: HomeAssistant, config_entry_hub, mqtt_mock: MqttMockHAClient
+) -> None:
+    """Test DROP switches for hubs."""
+    config_entry_hub.add_to_hass(hass)
+    assert await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    water_supply_switch_name = "switch.hub_drop_1_c0ffee_water_supply"
+    hass.states.async_set(water_supply_switch_name, STATE_UNKNOWN)
+    bypass_switch_name = "switch.hub_drop_1_c0ffee_treatment_bypass"
+    hass.states.async_set(bypass_switch_name, STATE_UNKNOWN)
+
+    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()
+
+    water_supply_switch = hass.states.get(water_supply_switch_name)
+    assert water_supply_switch
+    assert water_supply_switch.state == STATE_ON
+
+    # Test switch turn off method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: water_supply_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the hub
+    async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
+    await hass.async_block_till_done()
+
+    water_supply_switch = hass.states.get(water_supply_switch_name)
+    assert water_supply_switch
+    assert water_supply_switch.state == STATE_OFF
+
+    # Test switch turn on method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_ON,
+        {ATTR_ENTITY_ID: water_supply_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the hub
+    async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB)
+    await hass.async_block_till_done()
+
+    water_supply_switch = hass.states.get(water_supply_switch_name)
+    assert water_supply_switch
+    assert water_supply_switch.state == STATE_ON
+
+    bypass_switch = hass.states.get(bypass_switch_name)
+    assert bypass_switch
+    assert bypass_switch.state == STATE_OFF
+
+    # Test switch turn on method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_ON,
+        {ATTR_ENTITY_ID: bypass_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the device
+    async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB_RESET)
+    await hass.async_block_till_done()
+
+    bypass_switch = hass.states.get(bypass_switch_name)
+    assert bypass_switch
+    assert bypass_switch.state == STATE_ON
+
+    # Test switch turn off method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: bypass_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the device
+    async_fire_mqtt_message(hass, TEST_DATA_HUB_TOPIC, TEST_DATA_HUB)
+    await hass.async_block_till_done()
+
+    bypass_switch = hass.states.get(bypass_switch_name)
+    assert bypass_switch
+    assert bypass_switch.state == STATE_OFF
+
+
+async def test_switches_protection_valve(
+    hass: HomeAssistant, config_entry_protection_valve, mqtt_mock: MqttMockHAClient
+) -> None:
+    """Test DROP switches for protection valves."""
+    config_entry_protection_valve.add_to_hass(hass)
+    assert await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    async_fire_mqtt_message(
+        hass, TEST_DATA_PROTECTION_VALVE_TOPIC, TEST_DATA_PROTECTION_VALVE_RESET
+    )
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(
+        hass, TEST_DATA_PROTECTION_VALVE_TOPIC, TEST_DATA_PROTECTION_VALVE
+    )
+    await hass.async_block_till_done()
+
+    water_supply_switch_name = "switch.protection_valve_water_supply"
+    hass.states.async_set(water_supply_switch_name, STATE_UNKNOWN)
+
+    # Test switch turn off method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: water_supply_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the device
+    async_fire_mqtt_message(
+        hass, TEST_DATA_PROTECTION_VALVE_TOPIC, TEST_DATA_PROTECTION_VALVE_RESET
+    )
+    await hass.async_block_till_done()
+
+    water_supply_switch = hass.states.get(water_supply_switch_name)
+    assert water_supply_switch
+    assert water_supply_switch.state == STATE_OFF
+
+    # Test switch turn on method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_ON,
+        {ATTR_ENTITY_ID: water_supply_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the device
+    async_fire_mqtt_message(
+        hass, TEST_DATA_PROTECTION_VALVE_TOPIC, TEST_DATA_PROTECTION_VALVE
+    )
+    await hass.async_block_till_done()
+
+    water_supply_switch = hass.states.get(water_supply_switch_name)
+    assert water_supply_switch
+    assert water_supply_switch.state == STATE_ON
+
+
+async def test_switches_softener(
+    hass: HomeAssistant, config_entry_softener, mqtt_mock: MqttMockHAClient
+) -> None:
+    """Test DROP switches for softeners."""
+    config_entry_softener.add_to_hass(hass)
+    assert await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    async_fire_mqtt_message(hass, TEST_DATA_SOFTENER_TOPIC, TEST_DATA_SOFTENER_RESET)
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, TEST_DATA_SOFTENER_TOPIC, TEST_DATA_SOFTENER)
+    await hass.async_block_till_done()
+
+    bypass_switch_name = "switch.softener_treatment_bypass"
+    hass.states.async_set(bypass_switch_name, STATE_UNKNOWN)
+
+    # Test switch turn on method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_ON,
+        {ATTR_ENTITY_ID: bypass_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the device
+    async_fire_mqtt_message(hass, TEST_DATA_SOFTENER_TOPIC, TEST_DATA_SOFTENER_RESET)
+    await hass.async_block_till_done()
+
+    bypass_switch = hass.states.get(bypass_switch_name)
+    assert bypass_switch
+    assert bypass_switch.state == STATE_ON
+
+    # Test switch turn off method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: bypass_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the device
+    async_fire_mqtt_message(hass, TEST_DATA_SOFTENER_TOPIC, TEST_DATA_SOFTENER)
+    await hass.async_block_till_done()
+
+    bypass_switch = hass.states.get(bypass_switch_name)
+    assert bypass_switch
+    assert bypass_switch.state == STATE_OFF
+
+
+async def test_switches_filter(
+    hass: HomeAssistant, config_entry_filter, mqtt_mock: MqttMockHAClient
+) -> None:
+    """Test DROP switches for filters."""
+    config_entry_filter.add_to_hass(hass)
+    assert await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    async_fire_mqtt_message(hass, TEST_DATA_FILTER_TOPIC, TEST_DATA_FILTER_RESET)
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, TEST_DATA_FILTER_TOPIC, TEST_DATA_FILTER)
+    await hass.async_block_till_done()
+
+    bypass_switch_name = "switch.filter_treatment_bypass"
+    hass.states.async_set(bypass_switch_name, STATE_UNKNOWN)
+
+    # Test switch turn on method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_ON,
+        {ATTR_ENTITY_ID: bypass_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the device
+    async_fire_mqtt_message(hass, TEST_DATA_FILTER_TOPIC, TEST_DATA_FILTER_RESET)
+    await hass.async_block_till_done()
+
+    bypass_switch = hass.states.get(bypass_switch_name)
+    assert bypass_switch
+    assert bypass_switch.state == STATE_ON
+
+    # Test switch turn off method.
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: bypass_switch_name},
+        blocking=True,
+    )
+
+    # Simulate response from the device
+    async_fire_mqtt_message(hass, TEST_DATA_FILTER_TOPIC, TEST_DATA_FILTER)
+    await hass.async_block_till_done()
+
+    bypass_switch = hass.states.get(bypass_switch_name)
+    assert bypass_switch
+    assert bypass_switch.state == STATE_OFF