Use DataUpdateCoordinator for canary (#40691)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/components/canary/__init__.py

@@ -13,16 +13,16 @@ from homeassistant.const import CONF_PASSWORD, CONF_TIMEOUT, CONF_USERNAME
 from homeassistant.exceptions import ConfigEntryNotReady
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.typing import HomeAssistantType
-from homeassistant.util import Throttle
 
 from .const import (
     CONF_FFMPEG_ARGUMENTS,
-    DATA_CANARY,
+    DATA_COORDINATOR,
     DATA_UNDO_UPDATE_LISTENER,
     DEFAULT_FFMPEG_ARGUMENTS,
     DEFAULT_TIMEOUT,
     DOMAIN,
 )
+from .coordinator import CanaryDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -89,17 +89,21 @@ async def async_setup_entry(hass: HomeAssistantType, entry: ConfigEntry) -> bool
         hass.config_entries.async_update_entry(entry, options=options)
 
     try:
-        canary_data = await hass.async_add_executor_job(
+        canary_api = await hass.async_add_executor_job(_get_canary_api_instance, entry)
-            _get_canary_data_instance, entry
-        )
     except (ConnectTimeout, HTTPError) as error:
         _LOGGER.error("Unable to connect to Canary service: %s", str(error))
         raise ConfigEntryNotReady from error
 
+    coordinator = CanaryDataUpdateCoordinator(hass, api=canary_api)
+    await coordinator.async_refresh()
+
+    if not coordinator.last_update_success:
+        raise ConfigEntryNotReady
+
     undo_listener = entry.add_update_listener(_async_update_listener)
 
     hass.data[DOMAIN][entry.entry_id] = {
-        DATA_CANARY: canary_data,
+        DATA_COORDINATOR: coordinator,
         DATA_UNDO_UPDATE_LISTENER: undo_listener,
     }
 
@@ -134,77 +138,12 @@ async def _async_update_listener(hass: HomeAssistantType, entry: ConfigEntry) ->
     await hass.config_entries.async_reload(entry.entry_id)
 
 
-class CanaryData:
+def _get_canary_api_instance(entry: ConfigEntry) -> Api:
-    """Manages the data retrieved from Canary API."""
+    """Initialize a new instance of CanaryApi."""
-
-    def __init__(self, api: Api):
-        """Init the Canary data object."""
-        self._api = api
-        self._locations_by_id = {}
-        self._readings_by_device_id = {}
-
-    @Throttle(MIN_TIME_BETWEEN_UPDATES)
-    def update(self, **kwargs):
-        """Get the latest data from py-canary with a throttle."""
-        self._update(**kwargs)
-
-    def _update(self, **kwargs):
-        """Get the latest data from py-canary."""
-        for location in self._api.get_locations():
-            location_id = location.location_id
-
-            self._locations_by_id[location_id] = location
-
-            for device in location.devices:
-                if device.is_online:
-                    self._readings_by_device_id[
-                        device.device_id
-                    ] = self._api.get_latest_readings(device.device_id)
-
-    @property
-    def locations(self):
-        """Return a list of locations."""
-        return self._locations_by_id.values()
-
-    def get_location(self, location_id):
-        """Return a location based on location_id."""
-        return self._locations_by_id.get(location_id, [])
-
-    def get_readings(self, device_id):
-        """Return a list of readings based on device_id."""
-        return self._readings_by_device_id.get(device_id, [])
-
-    def get_reading(self, device_id, sensor_type):
-        """Return reading for device_id and sensor type."""
-        readings = self._readings_by_device_id.get(device_id, [])
-        return next(
-            (
-                reading.value
-                for reading in readings
-                if reading.sensor_type == sensor_type
-            ),
-            None,
-        )
-
-    def set_location_mode(self, location_id, mode_name, is_private=False):
-        """Set location mode."""
-        self._api.set_location_mode(location_id, mode_name, is_private)
-        self.update(no_throttle=True)
-
-    def get_live_stream_session(self, device):
-        """Return live stream session."""
-        return self._api.get_live_stream_session(device)
-
-
-def _get_canary_data_instance(entry: ConfigEntry) -> CanaryData:
-    """Initialize a new instance of CanaryData."""
     canary = Api(
         entry.data[CONF_USERNAME],
         entry.data[CONF_PASSWORD],
         entry.options.get(CONF_TIMEOUT, DEFAULT_TIMEOUT),
     )
 
-    canary_data = CanaryData(canary)
+    return canary
-    canary_data.update()
-
-    return canary_data

homeassistant/components/canary/alarm_control_panel.py

@@ -19,9 +19,10 @@ from homeassistant.const import (
 )
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.typing import HomeAssistantType
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import CanaryData
+from .const import DATA_COORDINATOR, DOMAIN
-from .const import DATA_CANARY, DOMAIN
+from .coordinator import CanaryDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -32,25 +33,35 @@ async def async_setup_entry(
     async_add_entities: Callable[[List[Entity], bool], None],
 ) -> None:
     """Set up Canary alarm control panels based on a config entry."""
-    data: CanaryData = hass.data[DOMAIN][entry.entry_id][DATA_CANARY]
+    coordinator: CanaryDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][
-    alarms = [CanaryAlarm(data, location.location_id) for location in data.locations]
+        DATA_COORDINATOR
+    ]
+    alarms = [
+        CanaryAlarm(coordinator, location)
+        for location_id, location in coordinator.data["locations"].items()
+    ]
 
     async_add_entities(alarms, True)
 
 
-class CanaryAlarm(AlarmControlPanelEntity):
+class CanaryAlarm(CoordinatorEntity, AlarmControlPanelEntity):
     """Representation of a Canary alarm control panel."""
 
-    def __init__(self, data, location_id):
+    def __init__(self, coordinator, location):
         """Initialize a Canary security camera."""
-        self._data = data
+        super().__init__(coordinator)
-        self._location_id = location_id
+        self._location_id = location.location_id
+        self._location_name = location.name
+
+    @property
+    def location(self):
+        """Return information about the location."""
+        return self.coordinator.data["locations"][self._location_id]
 
     @property
     def name(self):
         """Return the name of the alarm."""
-        location = self._data.get_location(self._location_id)
+        return self._location_name
-        return location.name
 
     @property
     def unique_id(self):
@@ -60,18 +71,17 @@ class CanaryAlarm(AlarmControlPanelEntity):
     @property
     def state(self):
         """Return the state of the device."""
-        location = self._data.get_location(self._location_id)
+        if self.location.is_private:
-
-        if location.is_private:
             return STATE_ALARM_DISARMED
 
-        mode = location.mode
+        mode = self.location.mode
         if mode.name == LOCATION_MODE_AWAY:
             return STATE_ALARM_ARMED_AWAY
         if mode.name == LOCATION_MODE_HOME:
             return STATE_ALARM_ARMED_HOME
         if mode.name == LOCATION_MODE_NIGHT:
             return STATE_ALARM_ARMED_NIGHT
+
         return None
 
     @property
@@ -82,26 +92,24 @@ class CanaryAlarm(AlarmControlPanelEntity):
     @property
     def device_state_attributes(self):
         """Return the state attributes."""
-        location = self._data.get_location(self._location_id)
+        return {"private": self.location.is_private}
-        return {"private": location.is_private}
 
     def alarm_disarm(self, code=None):
         """Send disarm command."""
-        location = self._data.get_location(self._location_id)
+        self.coordinator.canary.set_location_mode(
-        self._data.set_location_mode(self._location_id, location.mode.name, True)
+            self._location_id, self.location.mode.name, True
+        )
 
     def alarm_arm_home(self, code=None):
         """Send arm home command."""
-        self._data.set_location_mode(self._location_id, LOCATION_MODE_HOME)
+        self.coordinator.canary.set_location_mode(self._location_id, LOCATION_MODE_HOME)
 
     def alarm_arm_away(self, code=None):
         """Send arm away command."""
-        self._data.set_location_mode(self._location_id, LOCATION_MODE_AWAY)
+        self.coordinator.canary.set_location_mode(self._location_id, LOCATION_MODE_AWAY)
 
     def alarm_arm_night(self, code=None):
         """Send arm night command."""
-        self._data.set_location_mode(self._location_id, LOCATION_MODE_NIGHT)
+        self.coordinator.canary.set_location_mode(
-
+            self._location_id, LOCATION_MODE_NIGHT
-    def update(self):
+        )
-        """Get the latest state of the sensor."""
-        self._data.update()

homeassistant/components/canary/camera.py

@@ -15,17 +15,18 @@ from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.aiohttp_client import async_aiohttp_proxy_stream
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.typing import HomeAssistantType
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.util import Throttle
 
-from . import CanaryData
 from .const import (
     CONF_FFMPEG_ARGUMENTS,
-    DATA_CANARY,
+    DATA_COORDINATOR,
     DEFAULT_FFMPEG_ARGUMENTS,
     DEFAULT_TIMEOUT,
     DOMAIN,
     MANUFACTURER,
 )
+from .coordinator import CanaryDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -49,21 +50,22 @@ async def async_setup_entry(
     async_add_entities: Callable[[List[Entity], bool], None],
 ) -> None:
     """Set up Canary sensors based on a config entry."""
-    data: CanaryData = hass.data[DOMAIN][entry.entry_id][DATA_CANARY]
+    coordinator: CanaryDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][
-
+        DATA_COORDINATOR
+    ]
     ffmpeg_arguments = entry.options.get(
         CONF_FFMPEG_ARGUMENTS, DEFAULT_FFMPEG_ARGUMENTS
     )
     cameras = []
 
-    for location in data.locations:
+    for location_id, location in coordinator.data["locations"].items():
         for device in location.devices:
             if device.is_online:
                 cameras.append(
                     CanaryCamera(
                         hass,
-                        data,
+                        coordinator,
-                        location,
+                        location_id,
                         device,
                         DEFAULT_TIMEOUT,
                         ffmpeg_arguments,
@@ -73,17 +75,15 @@ async def async_setup_entry(
     async_add_entities(cameras, True)
 
 
-class CanaryCamera(Camera):
+class CanaryCamera(CoordinatorEntity, Camera):
     """An implementation of a Canary security camera."""
 
-    def __init__(self, hass, data, location, device, timeout, ffmpeg_args):
+    def __init__(self, hass, coordinator, location_id, device, timeout, ffmpeg_args):
         """Initialize a Canary security camera."""
-        super().__init__()
+        super().__init__(coordinator)
-
         self._ffmpeg = hass.data[DATA_FFMPEG]
         self._ffmpeg_arguments = ffmpeg_args
-        self._data = data
+        self._location_id = location_id
-        self._location = location
         self._device = device
         self._device_id = device.device_id
         self._device_name = device.name
@@ -91,6 +91,11 @@ class CanaryCamera(Camera):
         self._timeout = timeout
         self._live_stream_session = None
 
+    @property
+    def location(self):
+        """Return information about the location."""
+        return self.coordinator.data["locations"][self._location_id]
+
     @property
     def name(self):
         """Return the name of this device."""
@@ -114,12 +119,12 @@ class CanaryCamera(Camera):
     @property
     def is_recording(self):
         """Return true if the device is recording."""
-        return self._location.is_recording
+        return self.location.is_recording
 
     @property
     def motion_detection_enabled(self):
         """Return the camera motion detection status."""
-        return not self._location.is_recording
+        return not self.location.is_recording
 
     async def async_camera_image(self):
         """Return a still image response from the camera."""
@@ -159,4 +164,6 @@ class CanaryCamera(Camera):
     @Throttle(MIN_TIME_BETWEEN_SESSION_RENEW)
     def renew_live_stream_session(self):
         """Renew live stream session."""
-        self._live_stream_session = self._data.get_live_stream_session(self._device)
+        self._live_stream_session = self.coordinator.canary.get_live_stream_session(
+            self._device
+        )

homeassistant/components/canary/const.py

@@ -8,7 +8,7 @@ MANUFACTURER = "Canary Connect, Inc"
 CONF_FFMPEG_ARGUMENTS = "ffmpeg_arguments"
 
 # Data
-DATA_CANARY = "canary"
+DATA_COORDINATOR = "coordinator"
 DATA_UNDO_UPDATE_LISTENER = "undo_update_listener"
 
 # Defaults

homeassistant/components/canary/coordinator.py

@@ -0,0 +1,59 @@
+"""Provides the Canary DataUpdateCoordinator."""
+from datetime import timedelta
+import logging
+
+from async_timeout import timeout
+from canary.api import Api
+from requests import ConnectTimeout, HTTPError
+
+from homeassistant.helpers.typing import HomeAssistantType
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class CanaryDataUpdateCoordinator(DataUpdateCoordinator):
+    """Class to manage fetching Canary data."""
+
+    def __init__(self, hass: HomeAssistantType, *, api: Api):
+        """Initialize global Canary data updater."""
+        self.canary = api
+        update_interval = timedelta(seconds=30)
+
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=update_interval,
+        )
+
+    def _update_data(self) -> dict:
+        """Fetch data from Canary via sync functions."""
+        locations_by_id = {}
+        readings_by_device_id = {}
+
+        for location in self.canary.get_locations():
+            location_id = location.location_id
+            locations_by_id[location_id] = location
+
+            for device in location.devices:
+                if device.is_online:
+                    readings_by_device_id[
+                        device.device_id
+                    ] = self.canary.get_latest_readings(device.device_id)
+
+        return {
+            "locations": locations_by_id,
+            "readings": readings_by_device_id,
+        }
+
+    async def _async_update_data(self) -> dict:
+        """Fetch data from Canary."""
+
+        try:
+            async with timeout(15):
+                return await self.hass.async_add_executor_job(self._update_data)
+        except (ConnectTimeout, HTTPError) as error:
+            raise UpdateFailed(f"Invalid response from API: {error}") from error

homeassistant/components/canary/sensor.py

@@ -14,9 +14,10 @@ from homeassistant.const import (
 )
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.typing import HomeAssistantType
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import CanaryData
+from .const import DATA_COORDINATOR, DOMAIN, MANUFACTURER
-from .const import DATA_CANARY, DOMAIN, MANUFACTURER
+from .coordinator import CanaryDataUpdateCoordinator
 
 SENSOR_VALUE_PRECISION = 2
 ATTR_AIR_QUALITY = "air_quality"
@@ -49,37 +50,71 @@ async def async_setup_entry(
     async_add_entities: Callable[[List[Entity], bool], None],
 ) -> None:
     """Set up Canary sensors based on a config entry."""
-    data: CanaryData = hass.data[DOMAIN][entry.entry_id][DATA_CANARY]
+    coordinator: CanaryDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][
+        DATA_COORDINATOR
+    ]
     sensors = []
 
-    for location in data.locations:
+    for location in coordinator.data["locations"].values():
         for device in location.devices:
             if device.is_online:
                 device_type = device.device_type
                 for sensor_type in SENSOR_TYPES:
                     if device_type.get("name") in sensor_type[4]:
                         sensors.append(
-                            CanarySensor(data, sensor_type, location, device)
+                            CanarySensor(coordinator, sensor_type, location, device)
                         )
 
     async_add_entities(sensors, True)
 
 
-class CanarySensor(Entity):
+class CanarySensor(CoordinatorEntity, Entity):
     """Representation of a Canary sensor."""
 
-    def __init__(self, data, sensor_type, location, device):
+    def __init__(self, coordinator, sensor_type, location, device):
         """Initialize the sensor."""
-        self._data = data
+        super().__init__(coordinator)
         self._sensor_type = sensor_type
         self._device_id = device.device_id
         self._device_name = device.name
         self._device_type_name = device.device_type["name"]
-        self._sensor_value = None
 
         sensor_type_name = sensor_type[0].replace("_", " ").title()
         self._name = f"{location.name} {device.name} {sensor_type_name}"
 
+        canary_sensor_type = None
+        if self._sensor_type[0] == "air_quality":
+            canary_sensor_type = SensorType.AIR_QUALITY
+        elif self._sensor_type[0] == "temperature":
+            canary_sensor_type = SensorType.TEMPERATURE
+        elif self._sensor_type[0] == "humidity":
+            canary_sensor_type = SensorType.HUMIDITY
+        elif self._sensor_type[0] == "wifi":
+            canary_sensor_type = SensorType.WIFI
+        elif self._sensor_type[0] == "battery":
+            canary_sensor_type = SensorType.BATTERY
+
+        self._canary_type = canary_sensor_type
+
+    @property
+    def reading(self):
+        """Return the device sensor reading."""
+        readings = self.coordinator.data["readings"][self._device_id]
+
+        value = next(
+            (
+                reading.value
+                for reading in readings
+                if reading.sensor_type == self._canary_type
+            ),
+            None,
+        )
+
+        if value is not None:
+            return round(float(value), SENSOR_VALUE_PRECISION)
+
+        return None
+
     @property
     def name(self):
         """Return the name of the Canary sensor."""
@@ -88,7 +123,7 @@ class CanarySensor(Entity):
     @property
     def state(self):
         """Return the state of the sensor."""
-        return self._sensor_value
+        return self.reading
 
     @property
     def unique_id(self):
@@ -123,36 +158,17 @@ class CanarySensor(Entity):
     @property
     def device_state_attributes(self):
         """Return the state attributes."""
-        if self._sensor_type[0] == "air_quality" and self._sensor_value is not None:
+        reading = self.reading
+
+        if self._sensor_type[0] == "air_quality" and reading is not None:
             air_quality = None
-            if self._sensor_value <= 0.4:
+            if reading <= 0.4:
                 air_quality = STATE_AIR_QUALITY_VERY_ABNORMAL
-            elif self._sensor_value <= 0.59:
+            elif reading <= 0.59:
                 air_quality = STATE_AIR_QUALITY_ABNORMAL
-            elif self._sensor_value <= 1.0:
+            elif reading <= 1.0:
                 air_quality = STATE_AIR_QUALITY_NORMAL
 
             return {ATTR_AIR_QUALITY: air_quality}
 
         return None
-
-    def update(self):
-        """Get the latest state of the sensor."""
-        self._data.update()
-
-        canary_sensor_type = None
-        if self._sensor_type[0] == "air_quality":
-            canary_sensor_type = SensorType.AIR_QUALITY
-        elif self._sensor_type[0] == "temperature":
-            canary_sensor_type = SensorType.TEMPERATURE
-        elif self._sensor_type[0] == "humidity":
-            canary_sensor_type = SensorType.HUMIDITY
-        elif self._sensor_type[0] == "wifi":
-            canary_sensor_type = SensorType.WIFI
-        elif self._sensor_type[0] == "battery":
-            canary_sensor_type = SensorType.BATTERY
-
-        value = self._data.get_reading(self._device_id, canary_sensor_type)
-
-        if value is not None:
-            self._sensor_value = round(float(value), SENSOR_VALUE_PRECISION)

tests/components/canary/conftest.py

@@ -5,17 +5,12 @@ from pytest import fixture
 from tests.async_mock import MagicMock, patch
 
 
-def mock_canary_update(self, **kwargs):
-    """Get the latest data from py-canary."""
-    self._update(**kwargs)
-
-
 @fixture
 def canary(hass):
     """Mock the CanaryApi for easier testing."""
     with patch.object(Api, "login", return_value=True), patch(
-        "homeassistant.components.canary.CanaryData.update", mock_canary_update
+        "homeassistant.components.canary.Api"
-    ), patch("homeassistant.components.canary.Api") as mock_canary:
+    ) as mock_canary:
         instance = mock_canary.return_value = Api(
             "test-username",
             "test-password",

tests/components/canary/test_alarm_control_panel.py

@@ -137,7 +137,7 @@ async def test_alarm_control_panel_services(hass, canary) -> None:
         service_data={"entity_id": entity_id},
         blocking=True,
     )
-    instance.set_location_mode.assert_called_with(100, LOCATION_MODE_AWAY, False)
+    instance.set_location_mode.assert_called_with(100, LOCATION_MODE_AWAY)
 
     # test arm home
     await hass.services.async_call(
@@ -146,7 +146,7 @@ async def test_alarm_control_panel_services(hass, canary) -> None:
         service_data={"entity_id": entity_id},
         blocking=True,
     )
-    instance.set_location_mode.assert_called_with(100, LOCATION_MODE_HOME, False)
+    instance.set_location_mode.assert_called_with(100, LOCATION_MODE_HOME)
 
     # test arm night
     await hass.services.async_call(
@@ -155,7 +155,7 @@ async def test_alarm_control_panel_services(hass, canary) -> None:
         service_data={"entity_id": entity_id},
         blocking=True,
     )
-    instance.set_location_mode.assert_called_with(100, LOCATION_MODE_NIGHT, False)
+    instance.set_location_mode.assert_called_with(100, LOCATION_MODE_NIGHT)
 
     # test disarm
     await hass.services.async_call(

tests/components/canary/test_sensor.py

@@ -1,4 +1,6 @@
 """The tests for the Canary sensor platform."""
+from datetime import timedelta
+
 from homeassistant.components.canary.const import DOMAIN, MANUFACTURER
 from homeassistant.components.canary.sensor import (
     ATTR_AIR_QUALITY,
@@ -16,11 +18,12 @@ from homeassistant.const import (
     TEMP_CELSIUS,
 )
 from homeassistant.setup import async_setup_component
+from homeassistant.util.dt import utcnow
 
 from . import mock_device, mock_location, mock_reading
 
 from tests.async_mock import patch
-from tests.common import mock_device_registry, mock_registry
+from tests.common import async_fire_time_changed, mock_device_registry, mock_registry
 
 
 async def test_sensors_pro(hass, canary) -> None:
@@ -124,6 +127,8 @@ async def test_sensors_attributes_pro(hass, canary) -> None:
         mock_reading("air_quality", "0.4"),
     ]
 
+    future = utcnow() + timedelta(seconds=30)
+    async_fire_time_changed(hass, future)
     await hass.helpers.entity_component.async_update_entity(entity_id)
     await hass.async_block_till_done()
 
@@ -137,6 +142,8 @@ async def test_sensors_attributes_pro(hass, canary) -> None:
         mock_reading("air_quality", "1.0"),
     ]
 
+    future += timedelta(seconds=30)
+    async_fire_time_changed(hass, future)
     await hass.helpers.entity_component.async_update_entity(entity_id)
     await hass.async_block_till_done()