diff --git a/homeassistant/components/smarttub/__init__.py b/homeassistant/components/smarttub/__init__.py index 2b80c92510f..e8fc9989d38 100644 --- a/homeassistant/components/smarttub/__init__.py +++ b/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): diff --git a/homeassistant/components/smarttub/const.py b/homeassistant/components/smarttub/const.py index 99e7f21f86d..8292c5ef826 100644 --- a/homeassistant/components/smarttub/const.py +++ b/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 diff --git a/homeassistant/components/smarttub/controller.py b/homeassistant/components/smarttub/controller.py index ad40c94fbed..feb13066b44 100644 --- a/homeassistant/components/smarttub/controller.py +++ b/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.""" diff --git a/homeassistant/components/smarttub/manifest.json b/homeassistant/components/smarttub/manifest.json index 9735a3753b4..9360c59da8b 100644 --- a/homeassistant/components/smarttub/manifest.json +++ b/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" } diff --git a/homeassistant/components/smarttub/switch.py b/homeassistant/components/smarttub/switch.py new file mode 100644 index 00000000000..7e4c83f6feb --- /dev/null +++ b/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() diff --git a/requirements_all.txt b/requirements_all.txt index 9ceb668bc58..d8738946c19 100644 --- a/requirements_all.txt +++ b/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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index d37ee25cbbd..27f7475328e 100644 --- a/requirements_test_all.txt +++ b/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 diff --git a/tests/components/smarttub/conftest.py b/tests/components/smarttub/conftest.py index 265afcfc24c..fe1ca465f07 100644 --- a/tests/components/smarttub/conftest.py +++ b/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 diff --git a/tests/components/smarttub/test_sensor.py b/tests/components/smarttub/test_sensor.py index 8551ac3cccb..8e0cbf64abc 100644 --- a/tests/components/smarttub/test_sensor.py +++ b/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" diff --git a/tests/components/smarttub/test_switch.py b/tests/components/smarttub/test_switch.py new file mode 100644 index 00000000000..8750bf79747 --- /dev/null +++ b/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()