Add switch platform and pump entity to SmartTub (#46842)

homeassistant/components/smarttub/__init__.py

@@ -7,7 +7,7 @@ from .controller import SmartTubController
 
 _LOGGER = logging.getLogger(__name__)
 
-PLATFORMS = ["climate", "sensor"]
+PLATFORMS = ["climate", "sensor", "switch"]
 
 
 async def async_setup(hass, config):

homeassistant/components/smarttub/const.py

@@ -9,6 +9,7 @@ SMARTTUB_CONTROLLER = "smarttub_controller"
 SCAN_INTERVAL = 60
 
 POLLING_TIMEOUT = 10
+API_TIMEOUT = 5
 
 DEFAULT_MIN_TEMP = 18.5
 DEFAULT_MAX_TEMP = 40

homeassistant/components/smarttub/controller.py

@@ -86,7 +86,14 @@ class SmartTubController:
         return data
 
     async def _get_spa_data(self, spa):
-        return {"status": await spa.get_status()}
+        status, pumps = await asyncio.gather(
+            spa.get_status(),
+            spa.get_pumps(),
+        )
+        return {
+            "status": status,
+            "pumps": {pump.id: pump for pump in pumps},
+        }
 
     async def async_register_devices(self, entry):
         """Register devices with the device registry for all spas."""

homeassistant/components/smarttub/manifest.json

@@ -6,7 +6,7 @@
   "dependencies": [],
   "codeowners": ["@mdz"],
   "requirements": [
-      "python-smarttub==0.0.6"
+      "python-smarttub==0.0.12"
   ],
   "quality_scale": "platinum"
 }

homeassistant/components/smarttub/switch.py

@@ -0,0 +1,82 @@
+"""Platform for switch integration."""
+import logging
+
+import async_timeout
+from smarttub import SpaPump
+
+from homeassistant.components.switch import SwitchEntity
+
+from .const import API_TIMEOUT, DOMAIN, SMARTTUB_CONTROLLER
+from .entity import SmartTubEntity
+from .helpers import get_spa_name
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_entry(hass, entry, async_add_entities):
+    """Set up switch entities for the pumps on the tub."""
+
+    controller = hass.data[DOMAIN][entry.entry_id][SMARTTUB_CONTROLLER]
+
+    entities = [
+        SmartTubPump(controller.coordinator, pump)
+        for spa in controller.spas
+        for pump in await spa.get_pumps()
+    ]
+
+    async_add_entities(entities)
+
+
+class SmartTubPump(SmartTubEntity, SwitchEntity):
+    """A pump on a spa."""
+
+    def __init__(self, coordinator, pump: SpaPump):
+        """Initialize the entity."""
+        super().__init__(coordinator, pump.spa, "pump")
+        self.pump_id = pump.id
+        self.pump_type = pump.type
+
+    @property
+    def pump(self) -> SpaPump:
+        """Return the underlying SpaPump object for this entity."""
+        return self.coordinator.data[self.spa.id]["pumps"][self.pump_id]
+
+    @property
+    def unique_id(self) -> str:
+        """Return a unique ID for this pump entity."""
+        return f"{super().unique_id}-{self.pump_id}"
+
+    @property
+    def name(self) -> str:
+        """Return a name for this pump entity."""
+        spa_name = get_spa_name(self.spa)
+        if self.pump_type == SpaPump.PumpType.CIRCULATION:
+            return f"{spa_name} Circulation Pump"
+        if self.pump_type == SpaPump.PumpType.JET:
+            return f"{spa_name} Jet {self.pump_id}"
+        return f"{spa_name} pump {self.pump_id}"
+
+    @property
+    def is_on(self) -> bool:
+        """Return True if the pump is on."""
+        return self.pump.state != SpaPump.PumpState.OFF
+
+    async def async_turn_on(self, **kwargs) -> None:
+        """Turn the pump on."""
+
+        # the API only supports toggling
+        if not self.is_on:
+            await self.async_toggle()
+
+    async def async_turn_off(self, **kwargs) -> None:
+        """Turn the pump off."""
+
+        # the API only supports toggling
+        if self.is_on:
+            await self.async_toggle()
+
+    async def async_toggle(self, **kwargs) -> None:
+        """Toggle the pump on or off."""
+        async with async_timeout.timeout(API_TIMEOUT):
+            await self.pump.toggle()
+        await self.coordinator.async_request_refresh()

requirements_all.txt

@@ -1817,7 +1817,7 @@ python-qbittorrent==0.4.2
 python-ripple-api==0.0.3
 
 # homeassistant.components.smarttub
-python-smarttub==0.0.6
+python-smarttub==0.0.12
 
 # homeassistant.components.sochain
 python-sochain-api==0.0.2

requirements_test_all.txt

@@ -942,7 +942,7 @@ python-nest==4.1.0
 python-openzwave-mqtt[mqtt-client]==1.4.0
 
 # homeassistant.components.smarttub
-python-smarttub==0.0.6
+python-smarttub==0.0.12
 
 # homeassistant.components.songpal
 python-songpal==0.12

tests/components/smarttub/conftest.py

@@ -36,7 +36,7 @@ async def setup_component(hass):
 
 @pytest.fixture(name="spa")
 def mock_spa():
-    """Mock a SmartTub.Spa."""
+    """Mock a smarttub.Spa."""
 
     mock_spa = create_autospec(smarttub.Spa, instance=True)
     mock_spa.id = "mockspa1"
@@ -67,6 +67,27 @@ def mock_spa():
         "blowoutCycle": "INACTIVE",
         "cleanupCycle": "INACTIVE",
     }
+
+    mock_circulation_pump = create_autospec(smarttub.SpaPump, instance=True)
+    mock_circulation_pump.id = "CP"
+    mock_circulation_pump.spa = mock_spa
+    mock_circulation_pump.state = smarttub.SpaPump.PumpState.OFF
+    mock_circulation_pump.type = smarttub.SpaPump.PumpType.CIRCULATION
+
+    mock_jet_off = create_autospec(smarttub.SpaPump, instance=True)
+    mock_jet_off.id = "P1"
+    mock_jet_off.spa = mock_spa
+    mock_jet_off.state = smarttub.SpaPump.PumpState.OFF
+    mock_jet_off.type = smarttub.SpaPump.PumpType.JET
+
+    mock_jet_on = create_autospec(smarttub.SpaPump, instance=True)
+    mock_jet_on.id = "P2"
+    mock_jet_on.spa = mock_spa
+    mock_jet_on.state = smarttub.SpaPump.PumpState.HIGH
+    mock_jet_on.type = smarttub.SpaPump.PumpType.JET
+
+    mock_spa.get_pumps.return_value = [mock_circulation_pump, mock_jet_off, mock_jet_on]
+
     return mock_spa
 
 

tests/components/smarttub/test_sensor.py

@@ -3,7 +3,7 @@
 from . import trigger_update
 
 
-async def test_sensors(spa, setup_entry, hass, smarttub_api):
+async def test_sensors(spa, setup_entry, hass):
     """Test the sensors."""
 
     entity_id = f"sensor.{spa.brand}_{spa.model}_state"

tests/components/smarttub/test_switch.py

@@ -0,0 +1,38 @@
+"""Test the SmartTub switch platform."""
+
+from smarttub import SpaPump
+
+
+async def test_pumps(spa, setup_entry, hass):
+    """Test pump entities."""
+
+    for pump in spa.get_pumps.return_value:
+        if pump.type == SpaPump.PumpType.CIRCULATION:
+            entity_id = f"switch.{spa.brand}_{spa.model}_circulation_pump"
+        elif pump.type == SpaPump.PumpType.JET:
+            entity_id = f"switch.{spa.brand}_{spa.model}_jet_{pump.id.lower()}"
+        else:
+            raise NotImplementedError("Unknown pump type")
+
+        state = hass.states.get(entity_id)
+        assert state is not None
+        if pump.state == SpaPump.PumpState.OFF:
+            assert state.state == "off"
+
+            await hass.services.async_call(
+                "switch",
+                "turn_on",
+                {"entity_id": entity_id},
+                blocking=True,
+            )
+            pump.toggle.assert_called()
+        else:
+            assert state.state == "on"
+
+            await hass.services.async_call(
+                "switch",
+                "turn_off",
+                {"entity_id": entity_id},
+                blocking=True,
+            )
+            pump.toggle.assert_called()