Add support for real-time data from SimpliSafe (#31424)

* Add support for real-time data from SimpliSafe

* Updated requirements

* Linting

* Ensure dispatcher topic contains the domain

* Don't bother with a partial

* Websovket dataclass and other code review

* Ensure initial_event_to_use works with error

* Don't inline methods

* Don't abuse loop variable

* Simplify initial event retrieval

* Add connection lost and restored events

* Revert "Add connection lost and restored events"

This reverts commit e7ffe05.

* Make _on_disconnect a static method

* Code review comments

* Allow entities to opt out of REST and/or websocket API updates

* Revert "Allow entities to opt out of REST and/or websocket API updates"

This reverts commit 1989f2e.

* Code review comments

* Fix issues with events not triggering correct entities

* Bug fixes

homeassistant/components/simplisafe/__init__.py

@@ -1,10 +1,18 @@
 """Support for SimpliSafe alarm systems."""
 import asyncio
+from dataclasses import dataclass, field
+from datetime import datetime
 import logging
+from typing import Optional
 
 from simplipy import API
-from simplipy.errors import InvalidCredentialsError, SimplipyError
-from simplipy.system.v3 import VOLUME_HIGH, VOLUME_LOW, VOLUME_MEDIUM, VOLUME_OFF
+from simplipy.entity import EntityTypes
+from simplipy.errors import InvalidCredentialsError, SimplipyError, WebsocketError
+from simplipy.websocket import (
+    EVENT_LOCK_LOCKED,
+    EVENT_LOCK_UNLOCKED,
+    get_event_type_from_payload,
+)
 import voluptuous as vol
 
 from homeassistant.config_entries import SOURCE_IMPORT
@@ -21,36 +29,50 @@
     async_dispatcher_send,
 )
 from homeassistant.helpers.entity import Entity
-from homeassistant.helpers.event import async_track_time_interval
+from homeassistant.helpers.event import async_call_later, async_track_time_interval
 from homeassistant.helpers.service import (
     async_register_admin_service,
     verify_domain_control,
 )
+from homeassistant.util.dt import utc_from_timestamp
 
 from .config_flow import configured_instances
-from .const import DATA_CLIENT, DEFAULT_SCAN_INTERVAL, DOMAIN, TOPIC_UPDATE
+from .const import (
+    ATTR_ALARM_DURATION,
+    ATTR_ALARM_VOLUME,
+    ATTR_CHIME_VOLUME,
+    ATTR_ENTRY_DELAY_AWAY,
+    ATTR_ENTRY_DELAY_HOME,
+    ATTR_EXIT_DELAY_AWAY,
+    ATTR_EXIT_DELAY_HOME,
+    ATTR_LIGHT,
+    ATTR_VOICE_PROMPT_VOLUME,
+    DATA_CLIENT,
+    DEFAULT_SCAN_INTERVAL,
+    DOMAIN,
+    VOLUMES,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
 CONF_ACCOUNTS = "accounts"
 
 DATA_LISTENER = "listener"
+TOPIC_UPDATE = "simplisafe_update_data_{0}"
+
+DEFAULT_SOCKET_MIN_RETRY = 15
+DEFAULT_WATCHDOG_SECONDS = 5 * 60
+
+WEBSOCKET_EVENTS_REQUIRING_SERIAL = [EVENT_LOCK_LOCKED, EVENT_LOCK_UNLOCKED]
 
-ATTR_ALARM_DURATION = "alarm_duration"
-ATTR_ALARM_VOLUME = "alarm_volume"
-ATTR_CHIME_VOLUME = "chime_volume"
-ATTR_ENTRY_DELAY_AWAY = "entry_delay_away"
-ATTR_ENTRY_DELAY_HOME = "entry_delay_home"
-ATTR_EXIT_DELAY_AWAY = "exit_delay_away"
-ATTR_EXIT_DELAY_HOME = "exit_delay_home"
-ATTR_LIGHT = "light"
+ATTR_LAST_EVENT_INFO = "last_event_info"
+ATTR_LAST_EVENT_SENSOR_NAME = "last_event_sensor_name"
+ATTR_LAST_EVENT_SENSOR_TYPE = "last_event_sensor_type"
+ATTR_LAST_EVENT_TIMESTAMP = "last_event_timestamp"
 ATTR_PIN_LABEL = "label"
 ATTR_PIN_LABEL_OR_VALUE = "label_or_pin"
 ATTR_PIN_VALUE = "pin"
 ATTR_SYSTEM_ID = "system_id"
-ATTR_VOICE_PROMPT_VOLUME = "voice_prompt_volume"
-
-VOLUMES = [VOLUME_OFF, VOLUME_LOW, VOLUME_MEDIUM, VOLUME_HIGH]
 
 SERVICE_BASE_SCHEMA = vol.Schema({vol.Required(ATTR_SYSTEM_ID): cv.positive_int})
 
@@ -283,30 +305,167 @@ async def async_unload_entry(hass, entry):
     return True
 
 
+@dataclass(frozen=True)
+class SimpliSafeWebsocketEvent:
+    """Define a representation of a parsed websocket event."""
+
+    event_data: dict
+
+    changed_by: Optional[str] = field(init=False)
+    event_type: Optional[str] = field(init=False)
+    info: str = field(init=False)
+    sensor_name: str = field(init=False)
+    sensor_serial: str = field(init=False)
+    sensor_type: EntityTypes = field(init=False)
+    system_id: int = field(init=False)
+    timestamp: datetime = field(init=False)
+
+    def __post_init__(self):
+        """Initialize."""
+        object.__setattr__(self, "changed_by", self.event_data["pinName"])
+        object.__setattr__(
+            self, "event_type", get_event_type_from_payload(self.event_data)
+        )
+        object.__setattr__(self, "info", self.event_data["info"])
+        object.__setattr__(self, "sensor_name", self.event_data["sensorName"])
+        object.__setattr__(self, "sensor_serial", self.event_data["sensorSerial"])
+        try:
+            object.__setattr__(
+                self, "sensor_type", EntityTypes(self.event_data["sensorType"]).name
+            )
+        except ValueError:
+            _LOGGER.warning(
+                'Encountered unknown entity type: %s ("%s"). Please report it at'
+                "https://github.com/home-assistant/home-assistant/issues.",
+                self.event_data["sensorType"],
+                self.event_data["sensorName"],
+            )
+            object.__setattr__(self, "sensor_type", None)
+        object.__setattr__(self, "system_id", self.event_data["sid"])
+        object.__setattr__(
+            self, "timestamp", utc_from_timestamp(self.event_data["eventTimestamp"])
+        )
+
+
+class SimpliSafeWebsocket:
+    """Define a SimpliSafe websocket "manager" object."""
+
+    def __init__(self, hass, websocket):
+        """Initialize."""
+        self._hass = hass
+        self._websocket = websocket
+        self._websocket_reconnect_delay = DEFAULT_SOCKET_MIN_RETRY
+        self._websocket_reconnect_underway = False
+        self._websocket_watchdog_listener = None
+        self.last_events = {}
+
+    async def _async_attempt_websocket_connect(self):
+        """Attempt to connect to the websocket (retrying later on fail)."""
+        self._websocket_reconnect_underway = True
+
+        try:
+            await self._websocket.async_connect()
+        except WebsocketError as err:
+            _LOGGER.error("Error with the websocket connection: %s", err)
+            self._websocket_reconnect_delay = min(
+                2 * self._websocket_reconnect_delay, 480
+            )
+            async_call_later(
+                self._hass,
+                self._websocket_reconnect_delay,
+                self.async_websocket_connect,
+            )
+        else:
+            self._websocket_reconnect_delay = DEFAULT_SOCKET_MIN_RETRY
+            self._websocket_reconnect_underway = False
+
+    async def _async_websocket_reconnect(self, event_time):
+        """Forcibly disconnect from and reconnect to the websocket."""
+        _LOGGER.debug("Websocket watchdog expired; forcing socket reconnection")
+        await self.async_websocket_disconnect()
+        await self._async_attempt_websocket_connect()
+
+    def _on_connect(self):
+        """Define a handler to fire when the websocket is connected."""
+        _LOGGER.info("Connected to websocket")
+        _LOGGER.debug("Websocket watchdog starting")
+        if self._websocket_watchdog_listener is not None:
+            self._websocket_watchdog_listener()
+        self._websocket_watchdog_listener = async_call_later(
+            self._hass, DEFAULT_WATCHDOG_SECONDS, self._async_websocket_reconnect
+        )
+
+    @staticmethod
+    def _on_disconnect():
+        """Define a handler to fire when the websocket is disconnected."""
+        _LOGGER.info("Disconnected from websocket")
+
+    def _on_event(self, data):
+        """Define a handler to fire when a new SimpliSafe event arrives."""
+        event = SimpliSafeWebsocketEvent(data)
+        _LOGGER.debug("New websocket event: %s", event)
+        self.last_events[data["sid"]] = event
+        async_dispatcher_send(self._hass, TOPIC_UPDATE.format(data["sid"]))
+
+        _LOGGER.debug("Resetting websocket watchdog")
+        self._websocket_watchdog_listener()
+        self._websocket_watchdog_listener = async_call_later(
+            self._hass, DEFAULT_WATCHDOG_SECONDS, self._async_websocket_reconnect
+        )
+        self._websocket_reconnect_delay = DEFAULT_SOCKET_MIN_RETRY
+
+    async def async_websocket_connect(self):
+        """Register handlers and connect to the websocket."""
+        if self._websocket_reconnect_underway:
+            return
+
+        self._websocket.on_connect(self._on_connect)
+        self._websocket.on_disconnect(self._on_disconnect)
+        self._websocket.on_event(self._on_event)
+
+        await self._async_attempt_websocket_connect()
+
+    async def async_websocket_disconnect(self):
+        """Disconnect from the websocket."""
+        await self._websocket.async_disconnect()
+
+
 class SimpliSafe:
-    """Define a SimpliSafe API object."""
+    """Define a SimpliSafe data object."""
 
     def __init__(self, hass, api, config_entry):
         """Initialize."""
         self._api = api
         self._config_entry = config_entry
         self._emergency_refresh_token_used = False
         self._hass = hass
-        self.last_event_data = {}
+        self.initial_event_to_use = {}
         self.systems = None
+        self.websocket = SimpliSafeWebsocket(hass, api.websocket)
 
     async def async_init(self):
         """Initialize the data class."""
-        self.systems = await self._api.get_systems()
+        asyncio.create_task(self.websocket.async_websocket_connect())
 
-        # Register the base station for each system:
+        self.systems = await self._api.get_systems()
         for system in self.systems.values():
             self._hass.async_create_task(
                 async_register_base_station(
                     self._hass, system, self._config_entry.entry_id
                 )
             )
 
+            # Future events will come from the websocket, but since subscription to the
+            # websocket doesn't provide the most recent event, we grab it from the REST
+            # API to ensure event-related attributes aren't empty on startup:
+            try:
+                self.initial_event_to_use[
+                    system.system_id
+                ] = await system.get_latest_event()
+            except SimplipyError as err:
+                _LOGGER.error("Error while fetching initial event: %s", err)
+                self.initial_event_to_use[system.system_id] = {}
+
         async def refresh(event_time):
             """Refresh data from the SimpliSafe account."""
             await self.async_update()
@@ -323,7 +482,8 @@ async def async_update(self):
         async def update_system(system):
             """Update a system."""
             await system.update()
-            self.last_event_data[system.system_id] = await system.get_latest_event()
+            _LOGGER.debug('Updated REST API data for "%s"', system.address)
+            async_dispatcher_send(self._hass, TOPIC_UPDATE.format(system.system_id))
 
         tasks = [update_system(system) for system in self.systems.values()]
 
@@ -371,26 +531,41 @@ def cancel_tasks():
         if self._emergency_refresh_token_used:
             self._emergency_refresh_token_used = False
 
-        _LOGGER.debug("Updated data for all SimpliSafe systems")
-        async_dispatcher_send(self._hass, TOPIC_UPDATE)
-
 
 class SimpliSafeEntity(Entity):
     """Define a base SimpliSafe entity."""
 
-    def __init__(self, system, name, *, serial=None):
+    def __init__(self, simplisafe, system, name, *, serial=None):
         """Initialize."""
         self._async_unsub_dispatcher_connect = None
-        self._attrs = {ATTR_SYSTEM_ID: system.system_id}
+        self._last_processed_websocket_event = None
         self._name = name
         self._online = True
+        self._simplisafe = simplisafe
         self._system = system
+        self.websocket_events_to_listen_for = []
 
         if serial:
             self._serial = serial
         else:
             self._serial = system.serial
 
+        self._attrs = {
+            ATTR_LAST_EVENT_INFO: simplisafe.initial_event_to_use[system.system_id].get(
+                "info"
+            ),
+            ATTR_LAST_EVENT_SENSOR_NAME: simplisafe.initial_event_to_use[
+                system.system_id
+            ].get("sensorName"),
+            ATTR_LAST_EVENT_SENSOR_TYPE: simplisafe.initial_event_to_use[
+                system.system_id
+            ].get("sensorType"),
+            ATTR_LAST_EVENT_TIMESTAMP: simplisafe.initial_event_to_use[
+                system.system_id
+            ].get("eventTimestamp"),
+            ATTR_SYSTEM_ID: system.system_id,
+        }
+
     @property
     def available(self):
         """Return whether the entity is available."""
@@ -427,6 +602,36 @@ def unique_id(self):
         """Return the unique ID of the entity."""
         return self._serial
 
+    @callback
+    def _async_should_ignore_websocket_event(self, event):
+        """Return whether this entity should ignore a particular websocket event.
+
+        Note that we can't check for a final condition – whether the event belongs to
+        a particular entity, like a lock – because some events (like arming the system
+        from a keypad _or_ from the website) should impact the same entity.
+        """
+        # We've already processed this event:
+        if self._last_processed_websocket_event == event:
+            return True
+
+        # This is an event for a system other than the one this entity belongs to:
+        if event.system_id != self._system.system_id:
+            return True
+
+        # This isn't an event that this entity cares about:
+        if event.event_type not in self.websocket_events_to_listen_for:
+            return True
+
+        # This event is targeted at a specific entity whose serial number is different
+        # from this one's:
+        if (
+            event.event_type in WEBSOCKET_EVENTS_REQUIRING_SERIAL
+            and event.sensor_serial != self._serial
+        ):
+            return True
+
+        return False
+
     async def async_added_to_hass(self):
         """Register callbacks."""
 
@@ -436,9 +641,41 @@ def update():
             self.async_schedule_update_ha_state(True)
 
         self._async_unsub_dispatcher_connect = async_dispatcher_connect(
-            self.hass, TOPIC_UPDATE, update
+            self.hass, TOPIC_UPDATE.format(self._system.system_id), update
         )
 
+    async def async_update(self):
+        """Update the entity."""
+        self.async_update_from_rest_api()
+
+        last_websocket_event = self._simplisafe.websocket.last_events.get(
+            self._system.system_id
+        )
+
+        if self._async_should_ignore_websocket_event(last_websocket_event):
+            return
+
+        self._last_processed_websocket_event = last_websocket_event
+        self._attrs.update(
+            {
+                ATTR_LAST_EVENT_INFO: last_websocket_event.info,
+                ATTR_LAST_EVENT_SENSOR_NAME: last_websocket_event.sensor_name,
+                ATTR_LAST_EVENT_SENSOR_TYPE: last_websocket_event.sensor_type,
+                ATTR_LAST_EVENT_TIMESTAMP: last_websocket_event.timestamp,
+            }
+        )
+        self.async_update_from_websocket_event(last_websocket_event)
+
+    @callback
+    def async_update_from_rest_api(self):
+        """Update the entity with the provided REST API data."""
+        pass
+
+    @callback
+    def async_update_from_websocket_event(self, event):
+        """Update the entity with the provided websocket API data."""
+        pass
+
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect dispatcher listener when removed."""
         if self._async_unsub_dispatcher_connect: