Add entities for Balboa Spa pumps (#111245)

* Add entities for Balboa Spa pumps

* Fix fan tests and move client_update to __init__

* Ruff

---------

Co-authored-by: Jan-Philipp Benecke <github@bnck.me>

homeassistant/components/balboa/__init__.py

@@ -17,7 +17,7 @@ from .const import CONF_SYNC_TIME, DEFAULT_SYNC_TIME, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
-PLATFORMS = [Platform.BINARY_SENSOR, Platform.CLIMATE]
+PLATFORMS = [Platform.BINARY_SENSOR, Platform.CLIMATE, Platform.FAN]
 
 
 KEEP_ALIVE_INTERVAL = timedelta(minutes=1)

homeassistant/components/balboa/fan.py

@@ -0,0 +1,89 @@
+"""Support for Balboa Spa pumps."""
+from __future__ import annotations
+
+import math
+from typing import Any, cast
+
+from pybalboa import SpaClient, SpaControl
+from pybalboa.enums import OffOnState, UnknownState
+
+from homeassistant.components.fan import FanEntity, FanEntityFeature
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.util.percentage import (
+    percentage_to_ranged_value,
+    ranged_value_to_percentage,
+)
+
+from .const import DOMAIN
+from .entity import BalboaEntity
+
+
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
+    """Set up the spa's pumps."""
+    spa: SpaClient = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities(BalboaPumpFanEntity(control) for control in spa.pumps)
+
+
+class BalboaPumpFanEntity(BalboaEntity, FanEntity):
+    """Representation of a Balboa Spa pump fan entity."""
+
+    _attr_supported_features = FanEntityFeature.SET_SPEED
+    _attr_translation_key = "pump"
+
+    def __init__(self, control: SpaControl) -> None:
+        """Initialize a Balboa pump fan entity."""
+        super().__init__(control.client, control.name)
+        self._control = control
+        self._attr_translation_placeholders = {
+            "index": f"{cast(int, control.index) + 1}"
+        }
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn the pump off."""
+        await self._control.set_state(OffOnState.OFF)
+
+    async def async_turn_on(
+        self,
+        percentage: int | None = None,
+        preset_mode: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Turn the pump on (by default on max speed)."""
+        if percentage is None:
+            percentage = 100
+        await self.async_set_percentage(percentage)
+
+    async def async_set_percentage(self, percentage: int) -> None:
+        """Set the speed of the pump."""
+        if percentage > 0:
+            state = math.ceil(
+                percentage_to_ranged_value((1, self.speed_count), percentage)
+            )
+        else:
+            state = OffOnState.OFF
+        await self._control.set_state(state)
+
+    @property
+    def percentage(self) -> int | None:
+        """Return the speed of the pump."""
+        if self._control.state == UnknownState.UNKNOWN:
+            return None
+        if self._control.state == OffOnState.OFF:
+            return 0
+        return ranged_value_to_percentage((1, self.speed_count), self._control.state)
+
+    @property
+    def is_on(self) -> bool | None:
+        """Return true if the pump is running."""
+        if self._control.state == UnknownState.UNKNOWN:
+            return None
+        return self._control.state != OffOnState.OFF
+
+    @property
+    def speed_count(self) -> int:
+        """Return the number of different speed settings the pump supports."""
+        return int(max(self._control.options))

homeassistant/components/balboa/icons.json

@@ -19,6 +19,14 @@
           "on": "mdi:pump"
         }
       }
+    },
+    "fan": {
+      "pump": {
+        "default": "mdi:pump",
+        "state": {
+          "off": "mdi:pump-off"
+        }
+      }
     }
   }
 }

homeassistant/components/balboa/strings.json

@@ -52,6 +52,11 @@
           }
         }
       }
+    },
+    "fan": {
+      "pump": {
+        "name": "Pump {index}"
+      }
     }
   }
 }

tests/components/balboa/__init__.py

@@ -1,9 +1,11 @@
 """Test the Balboa Spa Client integration."""
 from __future__ import annotations
 
+from unittest.mock import MagicMock
+
 from homeassistant.components.balboa import CONF_SYNC_TIME, DOMAIN
 from homeassistant.const import CONF_HOST
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, State
 
 from tests.common import MockConfigEntry
 
@@ -19,3 +21,11 @@ async def init_integration(hass: HomeAssistant) -> MockConfigEntry:
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
     return entry
+
+
+async def client_update(hass: HomeAssistant, client: MagicMock, entity: str) -> State:
+    """Update the client."""
+    client.emit("")
+    await hass.async_block_till_done()
+    assert (state := hass.states.get(entity)) is not None
+    return state

tests/components/balboa/conftest.py

@@ -57,5 +57,6 @@ def client_fixture() -> Generator[MagicMock, None, None]:
         client.heat_mode.set_state = AsyncMock()
         client.heat_mode.options = list(HeatMode)[:2]
         client.heat_state = 2
+        client.pumps = []
 
         yield client

tests/components/balboa/test_climate.py

@@ -28,7 +28,7 @@ from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
 from homeassistant.core import HomeAssistant, State
 from homeassistant.exceptions import ServiceValidationError
 
-from . import init_integration
+from . import client_update, init_integration
 
 from tests.common import MockConfigEntry
 from tests.components.climate import common
@@ -149,7 +149,7 @@ async def test_spa_preset_modes(
         client.heat_mode.state = HeatMode[mode.upper()]
         await common.async_set_preset_mode(hass, mode, ENTITY_CLIMATE)
 
-        state = await _client_update(hass, client)
+        state = await client_update(hass, client, ENTITY_CLIMATE)
         assert state
         assert state.attributes[ATTR_PRESET_MODE] == mode
 
@@ -158,7 +158,7 @@ async def test_spa_preset_modes(
 
     # put it in RNR and test assertion
     client.heat_mode.state = HeatMode.READY_IN_REST
-    state = await _client_update(hass, client)
+    state = await client_update(hass, client, ENTITY_CLIMATE)
     assert state
     assert state.attributes[ATTR_PRESET_MODE] == "ready_in_rest"
 
@@ -199,19 +199,13 @@ async def test_spa_with_blower(hass: HomeAssistant, client: MagicMock) -> None:
 
 
 # Helpers
-async def _client_update(hass: HomeAssistant, client: MagicMock) -> State:
-    """Update the client."""
-    client.emit("")
-    await hass.async_block_till_done()
-    assert (state := hass.states.get(ENTITY_CLIMATE)) is not None
-    return state
 
 
 async def _patch_blower(hass: HomeAssistant, client: MagicMock, fan_mode: str) -> State:
     """Patch the blower state."""
     client.blowers[0].state = OffLowMediumHighState[fan_mode.upper()]
     await common.async_set_fan_mode(hass, fan_mode)
-    return await _client_update(hass, client)
+    return await client_update(hass, client, ENTITY_CLIMATE)
 
 
 async def _patch_spa_settemp(
@@ -223,7 +217,7 @@ async def _patch_spa_settemp(
     await common.async_set_temperature(
         hass, temperature=settemp, entity_id=ENTITY_CLIMATE
     )
-    return await _client_update(hass, client)
+    return await client_update(hass, client, ENTITY_CLIMATE)
 
 
 async def _patch_spa_heatmode(
@@ -232,7 +226,7 @@ async def _patch_spa_heatmode(
     """Patch the heatmode."""
     client.heat_mode.state = heat_mode
     await common.async_set_hvac_mode(hass, HVAC_SETTINGS[heat_mode], ENTITY_CLIMATE)
-    return await _client_update(hass, client)
+    return await client_update(hass, client, ENTITY_CLIMATE)
 
 
 async def _patch_spa_heatstate(
@@ -241,4 +235,4 @@ async def _patch_spa_heatstate(
     """Patch the heatmode."""
     client.heat_state = heat_state
     await common.async_set_hvac_mode(hass, HVAC_SETTINGS[heat_state], ENTITY_CLIMATE)
-    return await _client_update(hass, client)
+    return await client_update(hass, client, ENTITY_CLIMATE)

tests/components/balboa/test_fan.py

@@ -0,0 +1,82 @@
+"""Tests of the pump fan entity of the balboa integration."""
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+from pybalboa import SpaControl
+from pybalboa.enums import OffLowHighState, UnknownState
+import pytest
+
+from homeassistant.components.fan import ATTR_PERCENTAGE
+from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
+
+from . import client_update, init_integration
+
+from tests.components.fan import common
+
+ENTITY_FAN = "fan.fakespa_pump_1"
+
+
+@pytest.fixture
+def mock_pump(client: MagicMock):
+    """Return a mock pump."""
+    pump = MagicMock(SpaControl)
+
+    async def set_state(state: OffLowHighState):
+        pump.state = state
+
+    pump.client = client
+    pump.index = 0
+    pump.state = OffLowHighState.OFF
+    pump.set_state = set_state
+    pump.options = list(OffLowHighState)
+    client.pumps.append(pump)
+
+    return pump
+
+
+async def test_pump(hass: HomeAssistant, client: MagicMock, mock_pump) -> None:
+    """Test spa pump."""
+    await init_integration(hass)
+
+    # check if the initial state is off
+    state = hass.states.get(ENTITY_FAN)
+    assert state.state == STATE_OFF
+
+    # just call turn on, pump should be at full speed
+    await common.async_turn_on(hass, ENTITY_FAN)
+    state = await client_update(hass, client, ENTITY_FAN)
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_PERCENTAGE] == 100
+
+    # test setting percentage
+    await common.async_set_percentage(hass, ENTITY_FAN, 50)
+    state = await client_update(hass, client, ENTITY_FAN)
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_PERCENTAGE] == 50
+
+    # test calling turn off
+    await common.async_turn_off(hass, ENTITY_FAN)
+    state = await client_update(hass, client, ENTITY_FAN)
+    assert state.state == STATE_OFF
+
+    # test setting percentage to 0
+    await common.async_turn_on(hass, ENTITY_FAN)
+    await client_update(hass, client, ENTITY_FAN)
+
+    await common.async_set_percentage(hass, ENTITY_FAN, 0)
+    state = await client_update(hass, client, ENTITY_FAN)
+    assert state.state == STATE_OFF
+    assert state.attributes[ATTR_PERCENTAGE] == 0
+
+
+async def test_pump_unknown_state(
+    hass: HomeAssistant, client: MagicMock, mock_pump
+) -> None:
+    """Tests spa pump with unknown state."""
+    await init_integration(hass)
+
+    mock_pump.state = UnknownState.UNKNOWN
+    state = await client_update(hass, client, ENTITY_FAN)
+    assert state.state == STATE_UNKNOWN