diff --git a/homeassistant/components/smarttub/sensor.py b/homeassistant/components/smarttub/sensor.py index 320f288f36a..402fb87373f 100644 --- a/homeassistant/components/smarttub/sensor.py +++ b/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() diff --git a/tests/components/smarttub/conftest.py b/tests/components/smarttub/conftest.py index d1bd7c377e3..efba121822d 100644 --- a/tests/components/smarttub/conftest.py +++ b/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 diff --git a/tests/components/smarttub/test_sensor.py b/tests/components/smarttub/test_sensor.py index 7d62440295e..8551ac3cccb 100644 --- a/tests/components/smarttub/test_sensor.py +++ b/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"