Add support for SmartTub filtration cycles (#46868)

homeassistant/components/smarttub/climate.py

@@ -1,6 +1,8 @@
 """Platform for climate integration."""
 import logging
 
+from smarttub import Spa
+
 from homeassistant.components.climate import ClimateEntity
 from homeassistant.components.climate.const import (
     CURRENT_HVAC_HEAT,
@@ -38,9 +40,9 @@ class SmartTubThermostat(SmartTubEntity, ClimateEntity):
     """The target water temperature for the spa."""
 
     PRESET_MODES = {
-        "AUTO": PRESET_NONE,
+        Spa.HeatMode.AUTO: PRESET_NONE,
-        "ECO": PRESET_ECO,
+        Spa.HeatMode.ECONOMY: PRESET_ECO,
-        "DAY": PRESET_DAY,
+        Spa.HeatMode.DAY: PRESET_DAY,
     }
 
     HEAT_MODES = {v: k for k, v in PRESET_MODES.items()}
@@ -62,7 +64,7 @@ class SmartTubThermostat(SmartTubEntity, ClimateEntity):
     @property
     def hvac_action(self):
         """Return the current running hvac operation."""
-        return self.HVAC_ACTIONS.get(self.get_spa_status("heater"))
+        return self.HVAC_ACTIONS.get(self.spa_status.heater)
 
     @property
     def hvac_modes(self):
@@ -110,7 +112,7 @@ class SmartTubThermostat(SmartTubEntity, ClimateEntity):
     @property
     def preset_mode(self):
         """Return the current preset mode."""
-        return self.PRESET_MODES[self.get_spa_status("heatMode")]
+        return self.PRESET_MODES[self.spa_status.heat_mode]
 
     @property
     def preset_modes(self):
@@ -120,12 +122,12 @@ class SmartTubThermostat(SmartTubEntity, ClimateEntity):
     @property
     def current_temperature(self):
         """Return the current water temperature."""
-        return self.get_spa_status("water.temperature")
+        return self.spa_status.water.temperature
 
     @property
     def target_temperature(self):
         """Return the target water temperature."""
-        return self.get_spa_status("setTemperature")
+        return self.spa_status.set_temperature
 
     async def async_set_temperature(self, **kwargs):
         """Set new target temperature."""

homeassistant/components/smarttub/entity.py

@@ -52,18 +52,8 @@ class SmartTubEntity(CoordinatorEntity):
         spa_name = get_spa_name(self.spa)
         return f"{spa_name} {self._entity_type}"
 
-    def get_spa_status(self, path):
+    @property
-        """Retrieve a value from the data returned by Spa.get_status().
+    def spa_status(self) -> smarttub.SpaState:
+        """Retrieve the result of Spa.get_status()."""
 
-        Nested keys can be specified by a dotted path, e.g.
+        return self.coordinator.data[self.spa.id].get("status")
-        status['foo']['bar'] is 'foo.bar'.
-        """
-
-        status = self.coordinator.data[self.spa.id].get("status")
-        if status is None:
-            return None
-
-        for key in path.split("."):
-            status = status[key]
-
-        return status

homeassistant/components/smarttub/manifest.json

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

homeassistant/components/smarttub/sensor.py

@@ -1,4 +1,5 @@
 """Platform for sensor integration."""
+from enum import Enum
 import logging
 
 from .const import DOMAIN, SMARTTUB_CONTROLLER
@@ -6,6 +7,11 @@ from .entity import SmartTubEntity
 
 _LOGGER = logging.getLogger(__name__)
 
+ATTR_DURATION = "duration"
+ATTR_LAST_UPDATED = "last_updated"
+ATTR_MODE = "mode"
+ATTR_START_HOUR = "start_hour"
+
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up sensor entities for the sensors in the tub."""
@@ -18,15 +24,17 @@ async def async_setup_entry(hass, entry, async_add_entities):
             [
                 SmartTubSensor(controller.coordinator, spa, "State", "state"),
                 SmartTubSensor(
-                    controller.coordinator, spa, "Flow Switch", "flowSwitch"
+                    controller.coordinator, spa, "Flow Switch", "flow_switch"
                 ),
                 SmartTubSensor(controller.coordinator, spa, "Ozone", "ozone"),
                 SmartTubSensor(
-                    controller.coordinator, spa, "Blowout Cycle", "blowoutCycle"
+                    controller.coordinator, spa, "Blowout Cycle", "blowout_cycle"
                 ),
                 SmartTubSensor(
-                    controller.coordinator, spa, "Cleanup Cycle", "cleanupCycle"
+                    controller.coordinator, spa, "Cleanup Cycle", "cleanup_cycle"
                 ),
+                SmartTubPrimaryFiltrationCycle(controller.coordinator, spa),
+                SmartTubSecondaryFiltrationCycle(controller.coordinator, spa),
             ]
         )
 
@@ -36,17 +44,69 @@ async def async_setup_entry(hass, entry, async_add_entities):
 class SmartTubSensor(SmartTubEntity):
     """Generic and base class for SmartTub sensors."""
 
-    def __init__(self, coordinator, spa, sensor_name, spa_status_key):
+    def __init__(self, coordinator, spa, sensor_name, attr_name):
         """Initialize the entity."""
         super().__init__(coordinator, spa, sensor_name)
-        self._spa_status_key = spa_status_key
+        self._attr_name = attr_name
 
     @property
     def _state(self):
         """Retrieve the underlying state from the spa."""
-        return self.get_spa_status(self._spa_status_key)
+        return getattr(self.spa_status, self._attr_name)
 
     @property
     def state(self) -> str:
         """Return the current state of the sensor."""
+        if isinstance(self._state, Enum):
+            return self._state.name.lower()
         return self._state.lower()
+
+
+class SmartTubPrimaryFiltrationCycle(SmartTubSensor):
+    """The primary filtration cycle."""
+
+    def __init__(self, coordinator, spa):
+        """Initialize the entity."""
+        super().__init__(
+            coordinator, spa, "primary filtration cycle", "primary_filtration"
+        )
+
+    @property
+    def state(self) -> str:
+        """Return the current state of the sensor."""
+        return self._state.status.name.lower()
+
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        state = self._state
+        return {
+            ATTR_DURATION: state.duration,
+            ATTR_LAST_UPDATED: state.last_updated.isoformat(),
+            ATTR_MODE: state.mode.name.lower(),
+            ATTR_START_HOUR: state.start_hour,
+        }
+
+
+class SmartTubSecondaryFiltrationCycle(SmartTubSensor):
+    """The secondary filtration cycle."""
+
+    def __init__(self, coordinator, spa):
+        """Initialize the entity."""
+        super().__init__(
+            coordinator, spa, "Secondary Filtration Cycle", "secondary_filtration"
+        )
+
+    @property
+    def state(self) -> str:
+        """Return the current state of the sensor."""
+        return self._state.status.name.lower()
+
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        state = self._state
+        return {
+            ATTR_LAST_UPDATED: state.last_updated.isoformat(),
+            ATTR_MODE: state.mode.name.lower(),
+        }

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.12
+python-smarttub==0.0.17
 
 # 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.12
+python-smarttub==0.0.17
 
 # homeassistant.components.songpal
 python-songpal==0.12

tests/components/smarttub/conftest.py

@@ -42,7 +42,9 @@ def mock_spa():
     mock_spa.id = "mockspa1"
     mock_spa.brand = "mockbrand1"
     mock_spa.model = "mockmodel1"
-    mock_spa.get_status.return_value = {
+    mock_spa.get_status.return_value = smarttub.SpaState(
+        mock_spa,
+        **{
             "setTemperature": 39,
             "water": {"temperature": 38},
             "heater": "ON",
@@ -66,8 +68,8 @@ def mock_spa():
             "uv": "OFF",
             "blowoutCycle": "INACTIVE",
             "cleanupCycle": "INACTIVE",
-    }
+        },
-
+    )
     mock_circulation_pump = create_autospec(smarttub.SpaPump, instance=True)
     mock_circulation_pump.id = "CP"
     mock_circulation_pump.spa = mock_spa

tests/components/smarttub/test_climate.py

@@ -42,7 +42,7 @@ async def test_thermostat_update(spa, setup_entry, hass):
 
     assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_HEAT
 
-    spa.get_status.return_value["heater"] = "OFF"
+    spa.get_status.return_value.heater = "OFF"
     await trigger_update(hass)
     state = hass.states.get(entity_id)
 
@@ -83,9 +83,9 @@ async def test_thermostat_update(spa, setup_entry, hass):
         {ATTR_ENTITY_ID: entity_id, ATTR_PRESET_MODE: PRESET_ECO},
         blocking=True,
     )
-    spa.set_heat_mode.assert_called_with("ECO")
+    spa.set_heat_mode.assert_called_with(smarttub.Spa.HeatMode.ECONOMY)
 
-    spa.get_status.return_value["heatMode"] = "ECO"
+    spa.get_status.return_value.heat_mode = smarttub.Spa.HeatMode.ECONOMY
     await trigger_update(hass)
     state = hass.states.get(entity_id)
     assert state.attributes.get(ATTR_PRESET_MODE) == PRESET_ECO

tests/components/smarttub/test_sensor.py

@@ -11,7 +11,7 @@ async def test_sensors(spa, setup_entry, hass):
     assert state is not None
     assert state.state == "normal"
 
-    spa.get_status.return_value["state"] = "BAD"
+    spa.get_status.return_value.state = "BAD"
     await trigger_update(hass)
     state = hass.states.get(entity_id)
     assert state is not None
@@ -36,3 +36,19 @@ async def test_sensors(spa, setup_entry, hass):
     state = hass.states.get(entity_id)
     assert state is not None
     assert state.state == "inactive"
+
+    entity_id = f"sensor.{spa.brand}_{spa.model}_primary_filtration_cycle"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "inactive"
+    assert state.attributes["duration"] == 4
+    assert state.attributes["last_updated"] is not None
+    assert state.attributes["mode"] == "normal"
+    assert state.attributes["start_hour"] == 2
+
+    entity_id = f"sensor.{spa.brand}_{spa.model}_secondary_filtration_cycle"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "inactive"
+    assert state.attributes["last_updated"] is not None
+    assert state.attributes["mode"] == "away"