Add entity services to the Flo integration (#38287)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

homeassistant/components/flo/__init__.py

@@ -12,7 +12,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DOMAIN
+from .const import CLIENT, DOMAIN
 from .device import FloDeviceDataUpdateCoordinator
 
 CONFIG_SCHEMA = vol.Schema({DOMAIN: vol.Schema({})}, extra=vol.ALLOW_EXTRA)
@@ -30,10 +30,10 @@ async def async_setup(hass: HomeAssistant, config: dict):
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     """Set up flo from a config entry."""
-    hass.data[DOMAIN][entry.entry_id] = {}
     session = async_get_clientsession(hass)
+    hass.data[DOMAIN][entry.entry_id] = {}
     try:
-        hass.data[DOMAIN][entry.entry_id]["client"] = client = await async_get_api(
+        hass.data[DOMAIN][entry.entry_id][CLIENT] = client = await async_get_api(
             entry.data[CONF_USERNAME], entry.data[CONF_PASSWORD], session=session
         )
     except RequestError as err:
@@ -43,7 +43,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
 
     _LOGGER.debug("Flo user information with locations: %s", user_info)
 
-    hass.data[DOMAIN]["devices"] = devices = [
+    hass.data[DOMAIN][entry.entry_id]["devices"] = devices = [
         FloDeviceDataUpdateCoordinator(hass, client, location["id"], device["id"])
         for location in user_info["locations"]
         for device in location["devices"]

homeassistant/components/flo/binary_sensor.py

@@ -11,12 +11,12 @@ from .const import DOMAIN as FLO_DOMAIN
 from .device import FloDeviceDataUpdateCoordinator
 from .entity import FloEntity
 
-DEPENDENCIES = ["flo"]
-
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Flo sensors from config entry."""
-    devices: List[FloDeviceDataUpdateCoordinator] = hass.data[FLO_DOMAIN]["devices"]
+    devices: List[FloDeviceDataUpdateCoordinator] = hass.data[FLO_DOMAIN][
+        config_entry.entry_id
+    ]["devices"]
     entities = [FloPendingAlertsBinarySensor(device) for device in devices]
     async_add_entities(entities)
 

homeassistant/components/flo/const.py

@@ -1,3 +1,7 @@
 """Constants for the flo integration."""
-
+CLIENT = "client"
 DOMAIN = "flo"
+FLO_HOME = "home"
+FLO_AWAY = "away"
+FLO_SLEEP = "sleep"
+FLO_MODES = [FLO_HOME, FLO_AWAY, FLO_SLEEP]

homeassistant/components/flo/device.py

@@ -172,6 +172,24 @@ class FloDeviceDataUpdateCoordinator(DataUpdateCoordinator):
         """Return the target valve state for the device."""
         return self._device_information["valve"]["target"]
 
+    async def async_set_mode_home(self):
+        """Set the Flo location to home mode."""
+        await self.api_client.location.set_mode_home(self._flo_location_id)
+
+    async def async_set_mode_away(self):
+        """Set the Flo location to away mode."""
+        await self.api_client.location.set_mode_away(self._flo_location_id)
+
+    async def async_set_mode_sleep(self, sleep_minutes, revert_to_mode):
+        """Set the Flo location to sleep mode."""
+        await self.api_client.location.set_mode_sleep(
+            self._flo_location_id, sleep_minutes, revert_to_mode
+        )
+
+    async def async_run_health_test(self):
+        """Run a Flo device health test."""
+        await self.api_client.device.run_health_test(self._flo_device_id)
+
     async def _update_device(self, *_) -> None:
         """Update the device information from the API."""
         self._device_information = await self.api_client.device.get_info(

homeassistant/components/flo/manifest.json

@@ -3,7 +3,7 @@
   "name": "Flo",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/flo",
-  "requirements": ["aioflo==0.4.0"],
+  "requirements": ["aioflo==0.4.1"],
   "ssdp": [],
   "zeroconf": [],
   "homekit": {},

homeassistant/components/flo/sensor.py

@@ -15,8 +15,6 @@ from .const import DOMAIN as FLO_DOMAIN
 from .device import FloDeviceDataUpdateCoordinator
 from .entity import FloEntity
 
-DEPENDENCIES = ["flo"]
-
 WATER_ICON = "mdi:water"
 GAUGE_ICON = "mdi:gauge"
 NAME_DAILY_USAGE = "Today's Water Usage"
@@ -28,7 +26,9 @@ NAME_WATER_PRESSURE = "Water Pressure"
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Flo sensors from config entry."""
-    devices: List[FloDeviceDataUpdateCoordinator] = hass.data[FLO_DOMAIN]["devices"]
+    devices: List[FloDeviceDataUpdateCoordinator] = hass.data[FLO_DOMAIN][
+        config_entry.entry_id
+    ]["devices"]
     entities = []
     entities.extend([FloDailyUsageSensor(device) for device in devices])
     entities.extend([FloSystemModeSensor(device) for device in devices])

homeassistant/components/flo/services.yaml

@@ -0,0 +1,32 @@
+# Describes the format for available Flo services
+
+set_sleep_mode:
+  description: Set the location into sleep mode.
+  fields:
+    entity_id:
+      description: Flo switch entity id
+      example: "switch.shutoff_valve"
+    sleep_minutes:
+      description: The time to sleep in minutes.
+      example: 120
+    revert_to_mode:
+      description: The mode to revert to after sleep_minutes has elapsed.
+      example: "home"
+set_away_mode:
+  description: Set the location into away mode.
+  fields:
+    entity_id:
+      description: Flo switch entity id
+      example: "switch.shutoff_valve"
+set_home_mode:
+  description: Set the location into home mode.
+  fields:
+    entity_id:
+      description: Flo switch entity id
+      example: "switch.shutoff_valve"
+run_health_test:
+  description: Have the Flo device run a health test.
+  fields:
+    entity_id:
+      description: Flo switch entity id
+      example: "switch.shutoff_valve"

homeassistant/components/flo/switch.py

@@ -2,19 +2,54 @@
 
 from typing import List
 
+from aioflo.location import SLEEP_MINUTE_OPTIONS, SYSTEM_MODE_HOME, SYSTEM_REVERT_MODES
+import voluptuous as vol
+
 from homeassistant.components.switch import SwitchEntity
 from homeassistant.core import callback
+from homeassistant.helpers import entity_platform
 
 from .const import DOMAIN as FLO_DOMAIN
 from .device import FloDeviceDataUpdateCoordinator
 from .entity import FloEntity
 
+ATTR_REVERT_TO_MODE = "revert_to_mode"
+ATTR_SLEEP_MINUTES = "sleep_minutes"
+SERVICE_SET_SLEEP_MODE = "set_sleep_mode"
+SERVICE_SET_AWAY_MODE = "set_away_mode"
+SERVICE_SET_HOME_MODE = "set_home_mode"
+SERVICE_RUN_HEALTH_TEST = "run_health_test"
+
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Flo switches from config entry."""
-    devices: List[FloDeviceDataUpdateCoordinator] = hass.data[FLO_DOMAIN]["devices"]
+    devices: List[FloDeviceDataUpdateCoordinator] = hass.data[FLO_DOMAIN][
+        config_entry.entry_id
+    ]["devices"]
     async_add_entities([FloSwitch(device) for device in devices])
 
+    platform = entity_platform.current_platform.get()
+
+    platform.async_register_entity_service(
+        SERVICE_SET_AWAY_MODE, {}, "async_set_mode_away"
+    )
+    platform.async_register_entity_service(
+        SERVICE_SET_HOME_MODE, {}, "async_set_mode_home"
+    )
+    platform.async_register_entity_service(
+        SERVICE_RUN_HEALTH_TEST, {}, "async_run_health_test"
+    )
+    platform.async_register_entity_service(
+        SERVICE_SET_SLEEP_MODE,
+        {
+            vol.Required(ATTR_SLEEP_MINUTES, default=120): vol.In(SLEEP_MINUTE_OPTIONS),
+            vol.Required(ATTR_REVERT_TO_MODE, default=SYSTEM_MODE_HOME): vol.In(
+                SYSTEM_REVERT_MODES
+            ),
+        },
+        "async_set_mode_sleep",
+    )
+
 
 class FloSwitch(FloEntity, SwitchEntity):
     """Switch class for the Flo by Moen valve."""
@@ -57,3 +92,19 @@ class FloSwitch(FloEntity, SwitchEntity):
     async def async_added_to_hass(self):
         """When entity is added to hass."""
         self.async_on_remove(self._device.async_add_listener(self.async_update_state))
+
+    async def async_set_mode_home(self):
+        """Set the Flo location to home mode."""
+        await self._device.async_set_mode_home()
+
+    async def async_set_mode_away(self):
+        """Set the Flo location to away mode."""
+        await self._device.async_set_mode_away()
+
+    async def async_set_mode_sleep(self, sleep_minutes, revert_to_mode):
+        """Set the Flo location to sleep mode."""
+        await self._device.async_set_mode_sleep(sleep_minutes, revert_to_mode)
+
+    async def async_run_health_test(self):
+        """Run a Flo device health test."""
+        await self._device.async_run_health_test()

requirements_all.txt

@@ -163,7 +163,7 @@ aioeafm==0.1.2
 aioesphomeapi==2.6.1
 
 # homeassistant.components.flo
-aioflo==0.4.0
+aioflo==0.4.1
 
 # homeassistant.components.freebox
 aiofreepybox==0.0.8

requirements_test_all.txt

@@ -91,7 +91,7 @@ aioeafm==0.1.2
 aioesphomeapi==2.6.1
 
 # homeassistant.components.flo
-aioflo==0.4.0
+aioflo==0.4.1
 
 # homeassistant.components.freebox
 aiofreepybox==0.0.8

tests/components/flo/conftest.py

@@ -95,3 +95,40 @@ def aioclient_mock_fixture(aioclient_mock):
         headers={"Content-Type": "application/json"},
         json={"valve": {"target": "closed"}},
     )
+    # Mocks the health test call for flo.
+    aioclient_mock.post(
+        "https://api-gw.meetflo.com/api/v2/devices/98765/healthTest/run",
+        text=load_fixture("flo/user_info_expand_locations_response.json"),
+        status=200,
+        headers={"Content-Type": "application/json"},
+    )
+    # Mocks the health test call for flo.
+    aioclient_mock.post(
+        "https://api-gw.meetflo.com/api/v2/locations/mmnnoopp/systemMode",
+        text=load_fixture("flo/user_info_expand_locations_response.json"),
+        status=200,
+        headers={"Content-Type": "application/json"},
+        json={"systemMode": {"target": "home"}},
+    )
+    # Mocks the health test call for flo.
+    aioclient_mock.post(
+        "https://api-gw.meetflo.com/api/v2/locations/mmnnoopp/systemMode",
+        text=load_fixture("flo/user_info_expand_locations_response.json"),
+        status=200,
+        headers={"Content-Type": "application/json"},
+        json={"systemMode": {"target": "away"}},
+    )
+    # Mocks the health test call for flo.
+    aioclient_mock.post(
+        "https://api-gw.meetflo.com/api/v2/locations/mmnnoopp/systemMode",
+        text=load_fixture("flo/user_info_expand_locations_response.json"),
+        status=200,
+        headers={"Content-Type": "application/json"},
+        json={
+            "systemMode": {
+                "target": "sleep",
+                "revertMinutes": 120,
+                "revertMode": "home",
+            }
+        },
+    )

tests/components/flo/test_binary_sensor.py

@@ -19,7 +19,7 @@ async def test_binary_sensors(hass, config_entry, aioclient_mock_fixture):
     )
     await hass.async_block_till_done()
 
-    assert len(hass.data[FLO_DOMAIN]["devices"]) == 1
+    assert len(hass.data[FLO_DOMAIN][config_entry.entry_id]["devices"]) == 1
 
     # we should have 6 entities for the device
     state = hass.states.get("binary_sensor.pending_system_alerts")

tests/components/flo/test_device.py

@@ -19,9 +19,11 @@ async def test_device(hass, config_entry, aioclient_mock_fixture, aioclient_mock
         hass, FLO_DOMAIN, {CONF_USERNAME: TEST_USER_ID, CONF_PASSWORD: TEST_PASSWORD}
     )
     await hass.async_block_till_done()
-    assert len(hass.data[FLO_DOMAIN]["devices"]) == 1
+    assert len(hass.data[FLO_DOMAIN][config_entry.entry_id]["devices"]) == 1
 
-    device: FloDeviceDataUpdateCoordinator = hass.data[FLO_DOMAIN]["devices"][0]
+    device: FloDeviceDataUpdateCoordinator = hass.data[FLO_DOMAIN][
+        config_entry.entry_id
+    ]["devices"][0]
     assert device.api_client is not None
     assert device.available
     assert device.consumption_today == 3.674

tests/components/flo/test_init.py

@@ -13,4 +13,6 @@ async def test_setup_entry(hass, config_entry, aioclient_mock_fixture):
         hass, FLO_DOMAIN, {CONF_USERNAME: TEST_USER_ID, CONF_PASSWORD: TEST_PASSWORD}
     )
     await hass.async_block_till_done()
-    assert len(hass.data[FLO_DOMAIN]["devices"]) == 1
+    assert len(hass.data[FLO_DOMAIN][config_entry.entry_id]["devices"]) == 1
+
+    assert await hass.config_entries.async_unload(config_entry.entry_id)

tests/components/flo/test_sensor.py

@@ -14,7 +14,7 @@ async def test_sensors(hass, config_entry, aioclient_mock_fixture):
     )
     await hass.async_block_till_done()
 
-    assert len(hass.data[FLO_DOMAIN]["devices"]) == 1
+    assert len(hass.data[FLO_DOMAIN][config_entry.entry_id]["devices"]) == 1
 
     # we should have 5 entities for the device
     assert hass.states.get("sensor.current_system_mode").state == "home"
@@ -34,7 +34,7 @@ async def test_manual_update_entity(
     )
     await hass.async_block_till_done()
 
-    assert len(hass.data[FLO_DOMAIN]["devices"]) == 1
+    assert len(hass.data[FLO_DOMAIN][config_entry.entry_id]["devices"]) == 1
 
     await async_setup_component(hass, "homeassistant", {})
 

tests/components/flo/test_services.py

@@ -0,0 +1,69 @@
+"""Test the services for the Flo by Moen integration."""
+from homeassistant.components.flo.const import DOMAIN as FLO_DOMAIN
+from homeassistant.components.flo.switch import (
+    ATTR_REVERT_TO_MODE,
+    ATTR_SLEEP_MINUTES,
+    SERVICE_RUN_HEALTH_TEST,
+    SERVICE_SET_AWAY_MODE,
+    SERVICE_SET_HOME_MODE,
+    SERVICE_SET_SLEEP_MODE,
+    SYSTEM_MODE_HOME,
+)
+from homeassistant.const import ATTR_ENTITY_ID, CONF_PASSWORD, CONF_USERNAME
+from homeassistant.setup import async_setup_component
+
+from .common import TEST_PASSWORD, TEST_USER_ID
+
+SWITCH_ENTITY_ID = "switch.shutoff_valve"
+
+
+async def test_services(hass, config_entry, aioclient_mock_fixture, aioclient_mock):
+    """Test Flo services."""
+    config_entry.add_to_hass(hass)
+    assert await async_setup_component(
+        hass, FLO_DOMAIN, {CONF_USERNAME: TEST_USER_ID, CONF_PASSWORD: TEST_PASSWORD}
+    )
+    await hass.async_block_till_done()
+
+    assert len(hass.data[FLO_DOMAIN][config_entry.entry_id]["devices"]) == 1
+    assert aioclient_mock.call_count == 4
+
+    await hass.services.async_call(
+        FLO_DOMAIN,
+        SERVICE_RUN_HEALTH_TEST,
+        {ATTR_ENTITY_ID: SWITCH_ENTITY_ID},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert aioclient_mock.call_count == 5
+
+    await hass.services.async_call(
+        FLO_DOMAIN,
+        SERVICE_SET_AWAY_MODE,
+        {ATTR_ENTITY_ID: SWITCH_ENTITY_ID},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert aioclient_mock.call_count == 6
+
+    await hass.services.async_call(
+        FLO_DOMAIN,
+        SERVICE_SET_HOME_MODE,
+        {ATTR_ENTITY_ID: SWITCH_ENTITY_ID},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert aioclient_mock.call_count == 7
+
+    await hass.services.async_call(
+        FLO_DOMAIN,
+        SERVICE_SET_SLEEP_MODE,
+        {
+            ATTR_ENTITY_ID: SWITCH_ENTITY_ID,
+            ATTR_REVERT_TO_MODE: SYSTEM_MODE_HOME,
+            ATTR_SLEEP_MINUTES: 120,
+        },
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert aioclient_mock.call_count == 8

tests/components/flo/test_switch.py

@@ -15,7 +15,7 @@ async def test_valve_switches(hass, config_entry, aioclient_mock_fixture):
     )
     await hass.async_block_till_done()
 
-    assert len(hass.data[FLO_DOMAIN]["devices"]) == 1
+    assert len(hass.data[FLO_DOMAIN][config_entry.entry_id]["devices"]) == 1
 
     entity_id = "switch.shutoff_valve"
     assert hass.states.get(entity_id).state == STATE_ON