Add services to SmartTub for changing filtration settings

homeassistant/components/smarttub/entity.py

@@ -1,4 +1,4 @@
-"""SmartTub integration."""
+"""Base classes for SmartTub entities."""
 import logging
 
 import smarttub

homeassistant/components/smarttub/sensor.py

@@ -2,6 +2,11 @@
 from enum import Enum
 import logging
 
+import smarttub
+import voluptuous as vol
+
+from homeassistant.helpers import config_validation as cv, entity_platform
+
 from .const import DOMAIN, SMARTTUB_CONTROLLER
 from .entity import SmartTubSensorBase
 
@@ -14,6 +19,25 @@
 # the hour of the day at which to start the cycle (0-23)
 ATTR_START_HOUR = "start_hour"
 
+SET_PRIMARY_FILTRATION_SCHEMA = vol.All(
+    cv.has_at_least_one_key(ATTR_DURATION, ATTR_START_HOUR),
+    cv.make_entity_service_schema(
+        {
+            vol.Optional(ATTR_DURATION): vol.All(int, vol.Range(min=1, max=24)),
+            vol.Optional(ATTR_START_HOUR): vol.All(int, vol.Range(min=0, max=23)),
+        },
+    ),
+)
+
+SET_SECONDARY_FILTRATION_SCHEMA = {
+    vol.Required(ATTR_MODE): vol.In(
+        {
+            mode.name.lower()
+            for mode in smarttub.SpaSecondaryFiltrationCycle.SecondaryFiltrationMode
+        }
+    ),
+}
+
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up sensor entities for the sensors in the tub."""
@@ -43,6 +67,20 @@ async def async_setup_entry(hass, entry, async_add_entities):
 
     async_add_entities(entities)
 
+    platform = entity_platform.current_platform.get()
+
+    platform.async_register_entity_service(
+        "set_primary_filtration",
+        SET_PRIMARY_FILTRATION_SCHEMA,
+        "async_set_primary_filtration",
+    )
+
+    platform.async_register_entity_service(
+        "set_secondary_filtration",
+        SET_SECONDARY_FILTRATION_SCHEMA,
+        "async_set_secondary_filtration",
+    )
+
 
 class SmartTubSensor(SmartTubSensorBase):
     """Generic class for SmartTub status sensors."""
@@ -80,6 +118,13 @@ def device_state_attributes(self):
             ATTR_START_HOUR: state.start_hour,
         }
 
+    async def async_set_primary_filtration(self, **kwargs):
+        """Update primary filtration settings."""
+        await self._state.set(
+            duration=kwargs.get(ATTR_DURATION),
+            start_hour=kwargs.get(ATTR_START_HOUR),
+        )
+
 
 class SmartTubSecondaryFiltrationCycle(SmartTubSensor):
     """The secondary filtration cycle."""
@@ -103,3 +148,10 @@ def device_state_attributes(self):
             ATTR_CYCLE_LAST_UPDATED: state.last_updated.isoformat(),
             ATTR_MODE: state.mode.name.lower(),
         }
+
+    async def async_set_secondary_filtration(self, **kwargs):
+        """Update primary filtration settings."""
+        mode = smarttub.SpaSecondaryFiltrationCycle.SecondaryFiltrationMode[
+            kwargs[ATTR_MODE].upper()
+        ]
+        await self._state.set_mode(mode)

homeassistant/components/smarttub/services.yaml

@@ -0,0 +1,49 @@
+set_primary_filtration:
+  name: Update primary filtration settings
+  description: Updates the primary filtration settings
+  target:
+    entity:
+      integration: smarttub
+      domain: sensor
+  fields:
+    duration:
+      name: Duration
+      description: The desired duration of the primary filtration cycle
+      required: true
+      default: 8
+      selector:
+        number:
+          min: 1
+          max: 24
+          unit_of_measurement: "hours"
+          mode: slider
+      example: 8
+    start_hour:
+      description: The hour of the day at which to begin the primary filtration cycle
+      required: true
+      default: 0
+      example: 2
+      selector:
+        number:
+          min: 0
+          max: 23
+          unit_of_measurement: "hour"
+
+set_secondary_filtration:
+  name: Update secondary filtration settings
+  description: Updates the secondary filtration settings
+  target:
+    entity:
+      integration: smarttub
+      domain: sensor
+  fields:
+    mode:
+      description: The secondary filtration mode.
+      selector:
+        select:
+          options:
+            - "frequent"
+            - "infrequent"
+            - "away"
+      required: true
+      example: "frequent"

tests/components/smarttub/conftest.py

@@ -34,15 +34,11 @@ async def setup_component(hass):
     assert await async_setup_component(hass, DOMAIN, {}) is True
 
 
-@pytest.fixture(name="spa")
-def mock_spa():
-    """Mock a smarttub.Spa."""
+@pytest.fixture(name="spa_state")
+def mock_spa_state():
+    """Create a smarttub.SpaState with mocks."""
 
-    mock_spa = create_autospec(smarttub.Spa, instance=True)
-    mock_spa.id = "mockspa1"
-    mock_spa.brand = "mockbrand1"
-    mock_spa.model = "mockmodel1"
-    mock_spa.get_status.return_value = smarttub.SpaState(
+    state = smarttub.SpaState(
         mock_spa,
         **{
             "setTemperature": 39,
@@ -71,6 +67,26 @@ def mock_spa():
             "cleanupCycle": "INACTIVE",
         },
     )
+    state.primary_filtration.set = create_autospec(
+        smarttub.SpaPrimaryFiltrationCycle, instance=True
+    ).set
+    state.secondary_filtration.set_mode = create_autospec(
+        smarttub.SpaSecondaryFiltrationCycle, instance=True
+    ).set_mode
+    return state
+
+
+@pytest.fixture(name="spa")
+def mock_spa(spa_state):
+    """Mock a smarttub.Spa."""
+
+    mock_spa = create_autospec(smarttub.Spa, instance=True)
+    mock_spa.id = "mockspa1"
+    mock_spa.brand = "mockbrand1"
+    mock_spa.model = "mockmodel1"
+
+    mock_spa.get_status.return_value = spa_state
+
     mock_circulation_pump = create_autospec(smarttub.SpaPump, instance=True)
     mock_circulation_pump.id = "CP"
     mock_circulation_pump.spa = mock_spa

tests/components/smarttub/test_sensor.py

@@ -1,6 +1,7 @@
 """Test the SmartTub sensor platform."""
 
 import pytest
+import smarttub
 
 
 @pytest.mark.parametrize(
@@ -35,6 +36,16 @@ async def test_primary_filtration(spa, setup_entry, hass):
     assert state.attributes["mode"] == "normal"
     assert state.attributes["start_hour"] == 2
 
+    await hass.services.async_call(
+        "smarttub",
+        "set_primary_filtration",
+        {"entity_id": entity_id, "duration": 8, "start_hour": 1},
+        blocking=True,
+    )
+    spa.get_status.return_value.primary_filtration.set.assert_called_with(
+        duration=8, start_hour=1
+    )
+
 
 async def test_secondary_filtration(spa, setup_entry, hass):
     """Test the secondary filtration cycle sensor."""
@@ -45,3 +56,16 @@ async def test_secondary_filtration(spa, setup_entry, hass):
     assert state.state == "inactive"
     assert state.attributes["cycle_last_updated"] is not None
     assert state.attributes["mode"] == "away"
+
+    await hass.services.async_call(
+        "smarttub",
+        "set_secondary_filtration",
+        {
+            "entity_id": entity_id,
+            "mode": "frequent",
+        },
+        blocking=True,
+    )
+    spa.get_status.return_value.secondary_filtration.set_mode.assert_called_with(
+        mode=smarttub.SpaSecondaryFiltrationCycle.SecondaryFiltrationMode.FREQUENT
+    )