Add more sensors to SmartTub integration (#46839)

homeassistant/components/smarttub/sensor.py

@@ -8,23 +8,45 @@ _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(hass, entry, async_add_entities):
-    """Set up climate entity for the thermostat in the tub."""
+    """Set up sensor entities for the sensors in the tub."""
 
     controller = hass.data[DOMAIN][entry.entry_id][SMARTTUB_CONTROLLER]
 
-    entities = [SmartTubState(controller.coordinator, spa) for spa in controller.spas]
+    entities = []
+    for spa in controller.spas:
+        entities.extend(
+            [
+                SmartTubSensor(controller.coordinator, spa, "State", "state"),
+                SmartTubSensor(
+                    controller.coordinator, spa, "Flow Switch", "flowSwitch"
+                ),
+                SmartTubSensor(controller.coordinator, spa, "Ozone", "ozone"),
+                SmartTubSensor(
+                    controller.coordinator, spa, "Blowout Cycle", "blowoutCycle"
+                ),
+                SmartTubSensor(
+                    controller.coordinator, spa, "Cleanup Cycle", "cleanupCycle"
+                ),
+            ]
+        )
 
     async_add_entities(entities)
 
 
-class SmartTubState(SmartTubEntity):
-    """The state of the spa."""
+class SmartTubSensor(SmartTubEntity):
+    """Generic and base class for SmartTub sensors."""
 
-    def __init__(self, coordinator, spa):
+    def __init__(self, coordinator, spa, sensor_name, spa_status_key):
         """Initialize the entity."""
-        super().__init__(coordinator, spa, "state")
+        super().__init__(coordinator, spa, sensor_name)
+        self._spa_status_key = spa_status_key
+
+    @property
+    def _state(self):
+        """Retrieve the underlying state from the spa."""
+        return self.get_spa_status(self._spa_status_key)
 
     @property
     def state(self) -> str:
         """Return the current state of the sensor."""
-        return self.get_spa_status("state").lower()
+        return self._state.lower()

tests/components/smarttub/conftest.py

@@ -47,6 +47,24 @@ def mock_spa():
         "water": {"temperature": 38},
         "heater": "ON",
         "state": "NORMAL",
+        "primaryFiltration": {
+            "cycle": 1,
+            "duration": 4,
+            "lastUpdated": "2021-01-20T11:38:57.014Z",
+            "mode": "NORMAL",
+            "startHour": 2,
+            "status": "INACTIVE",
+        },
+        "secondaryFiltration": {
+            "lastUpdated": "2020-07-09T19:39:52.961Z",
+            "mode": "AWAY",
+            "status": "INACTIVE",
+        },
+        "flowSwitch": "OPEN",
+        "ozone": "OFF",
+        "uv": "OFF",
+        "blowoutCycle": "INACTIVE",
+        "cleanupCycle": "INACTIVE",
     }
     return mock_spa
 

tests/components/smarttub/test_sensor.py

@@ -3,8 +3,8 @@
 from . import trigger_update
 
 
-async def test_state_update(spa, setup_entry, hass, smarttub_api):
-    """Test the state entity."""
+async def test_sensors(spa, setup_entry, hass, smarttub_api):
+    """Test the sensors."""
 
     entity_id = f"sensor.{spa.brand}_{spa.model}_state"
     state = hass.states.get(entity_id)
@@ -16,3 +16,23 @@ async def test_state_update(spa, setup_entry, hass, smarttub_api):
     state = hass.states.get(entity_id)
     assert state is not None
     assert state.state == "bad"
+
+    entity_id = f"sensor.{spa.brand}_{spa.model}_flow_switch"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "open"
+
+    entity_id = f"sensor.{spa.brand}_{spa.model}_ozone"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "off"
+
+    entity_id = f"sensor.{spa.brand}_{spa.model}_blowout_cycle"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "inactive"
+
+    entity_id = f"sensor.{spa.brand}_{spa.model}_cleanup_cycle"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "inactive"