Add Response switch platform to FireServiceRota integration (#43700)

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

.coveragerc

@@ -266,6 +266,7 @@ omit =
     homeassistant/components/fireservicerota/binary_sensor.py
     homeassistant/components/fireservicerota/const.py
     homeassistant/components/fireservicerota/sensor.py
+    homeassistant/components/fireservicerota/switch.py
     homeassistant/components/firmata/__init__.py
     homeassistant/components/firmata/binary_sensor.py
     homeassistant/components/firmata/board.py

homeassistant/components/fireservicerota/__init__.py

@@ -13,6 +13,7 @@ from pyfireservicerota import (
 
 from homeassistant.components.binary_sensor import DOMAIN as BINARYSENSOR_DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
+from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntry
 from homeassistant.const import CONF_TOKEN, CONF_URL, CONF_USERNAME
 from homeassistant.core import HomeAssistant
@@ -25,7 +26,7 @@ MIN_TIME_BETWEEN_UPDATES = timedelta(seconds=60)
 
 _LOGGER = logging.getLogger(__name__)
 
-SUPPORTED_PLATFORMS = {SENSOR_DOMAIN, BINARYSENSOR_DOMAIN}
+SUPPORTED_PLATFORMS = {SENSOR_DOMAIN, BINARYSENSOR_DOMAIN, SWITCH_DOMAIN}
 
 
 async def async_setup(hass: HomeAssistant, config: dict) -> bool:
@@ -191,6 +192,7 @@ class FireServiceRotaClient:
 
         self.token_refresh_failure = False
         self.incident_id = None
+        self.on_duty = False
 
         self.fsr = FireServiceRota(base_url=self._url, token_info=self._tokens)
 
@@ -229,24 +231,29 @@ class FireServiceRotaClient:
             self.fsr.get_availability, str(self._hass.config.time_zone)
         )
 
+        self.on_duty = bool(data.get("available"))
+
         _LOGGER.debug("Updated availability data: %s", data)
         return data
 
     async def async_response_update(self) -> object:
         """Get the latest incident response data."""
-        data = self.websocket.incident_data()
-        if data is None or "id" not in data:
+
+        if not self.incident_id:
             return
 
-        self.incident_id = data("id")
-        _LOGGER.debug("Updating incident response data for id: %s", self.incident_id)
+        _LOGGER.debug("Updating response data for incident id %s", self.incident_id)
 
         return await self.update_call(self.fsr.get_incident_response, self.incident_id)
 
     async def async_set_response(self, value) -> None:
         """Set incident response status."""
+
+        if not self.incident_id:
+            return
+
         _LOGGER.debug(
-            "Setting incident response for incident '%s' to status '%s'",
+            "Setting incident response for incident id '%s' to state '%s'",
             self.incident_id,
             value,
         )

homeassistant/components/fireservicerota/binary_sensor.py

@@ -9,7 +9,7 @@ from homeassistant.helpers.update_coordinator import (
     DataUpdateCoordinator,
 )
 
-from .const import DATA_COORDINATOR, DOMAIN as FIRESERVICEROTA_DOMAIN
+from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN as FIRESERVICEROTA_DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -19,19 +19,22 @@ async def async_setup_entry(
 ) -> None:
     """Set up FireServiceRota binary sensor based on a config entry."""
 
+    client = hass.data[FIRESERVICEROTA_DOMAIN][entry.entry_id][DATA_CLIENT]
+
     coordinator: DataUpdateCoordinator = hass.data[FIRESERVICEROTA_DOMAIN][
         entry.entry_id
     ][DATA_COORDINATOR]
 
-    async_add_entities([ResponseBinarySensor(coordinator, entry)])
+    async_add_entities([ResponseBinarySensor(coordinator, client, entry)])
 
 
 class ResponseBinarySensor(CoordinatorEntity, BinarySensorEntity):
     """Representation of an FireServiceRota sensor."""
 
-    def __init__(self, coordinator: DataUpdateCoordinator, entry):
+    def __init__(self, coordinator: DataUpdateCoordinator, client, entry):
         """Initialize."""
         super().__init__(coordinator)
+        self._client = client
         self._unique_id = f"{entry.unique_id}_Duty"
 
         self._state = None
@@ -44,7 +47,10 @@ class ResponseBinarySensor(CoordinatorEntity, BinarySensorEntity):
     @property
     def icon(self) -> str:
         """Return the icon to use in the frontend."""
-        return "mdi:calendar"
+        if self._state:
+            return "mdi:calendar-check"
+
+        return "mdi:calendar-remove"
 
     @property
     def unique_id(self) -> str:
@@ -52,16 +58,10 @@ class ResponseBinarySensor(CoordinatorEntity, BinarySensorEntity):
         return self._unique_id
 
     @property
-    def is_on(self):
+    def is_on(self) -> bool:
         """Return the state of the binary sensor."""
-        if not self.coordinator.data:
-            return
 
-        data = self.coordinator.data
-        if "available" in data and data["available"]:
-            self._state = True
-        else:
-            self._state = False
+        self._state = self._client.on_duty
 
         _LOGGER.debug("Set state of entity 'Duty Binary Sensor' to '%s'", self._state)
         return self._state

homeassistant/components/fireservicerota/config_flow.py

@@ -5,7 +5,8 @@ import voluptuous as vol
 from homeassistant import config_entries
 from homeassistant.const import CONF_PASSWORD, CONF_TOKEN, CONF_URL, CONF_USERNAME
 
-from .const import DOMAIN, URL_LIST  # pylint: disable=unused-import
+# pylint: disable=relative-beyond-top-level
+from .const import DOMAIN, URL_LIST
 
 DATA_SCHEMA = vol.Schema(
     {

homeassistant/components/fireservicerota/sensor.py

@@ -7,6 +7,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import HomeAssistantType
 
+# pylint: disable=relative-beyond-top-level
 from .const import DATA_CLIENT, DOMAIN as FIRESERVICEROTA_DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
@@ -73,6 +74,7 @@ class IncidentsSensor(RestoreEntity):
             return attr
 
         for value in (
+            "id",
             "trigger",
             "created_at",
             "message_to_speech_url",
@@ -106,7 +108,9 @@ class IncidentsSensor(RestoreEntity):
         if state:
             self._state = state.state
             self._state_attributes = state.attributes
-            _LOGGER.debug("Restored entity 'Incidents' state to: %s", self._state)
+            if "id" in self._state_attributes:
+                self._client.incident_id = self._state_attributes["id"]
+            _LOGGER.debug("Restored entity 'Incidents' to: %s", self._state)
 
         self.async_on_remove(
             async_dispatcher_connect(
@@ -125,4 +129,6 @@ class IncidentsSensor(RestoreEntity):
 
         self._state = data["body"]
         self._state_attributes = data
+        if "id" in self._state_attributes:
+            self._client.incident_id = self._state_attributes["id"]
         self.async_write_ha_state()

homeassistant/components/fireservicerota/switch.py

@@ -0,0 +1,154 @@
+"""Switch platform for FireServiceRota integration."""
+import logging
+
+from homeassistant.components.switch import SwitchEntity
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import callback
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.typing import HomeAssistantType
+
+# pylint: disable=relative-beyond-top-level
+from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN as FIRESERVICEROTA_DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_entry(
+    hass: HomeAssistantType, entry: ConfigEntry, async_add_entities
+) -> None:
+    """Set up FireServiceRota switch based on a config entry."""
+    client = hass.data[FIRESERVICEROTA_DOMAIN][entry.entry_id][DATA_CLIENT]
+
+    coordinator = hass.data[FIRESERVICEROTA_DOMAIN][entry.entry_id][DATA_COORDINATOR]
+
+    async_add_entities([ResponseSwitch(coordinator, client, entry)])
+
+
+class ResponseSwitch(SwitchEntity):
+    """Representation of an FireServiceRota switch."""
+
+    def __init__(self, coordinator, client, entry):
+        """Initialize."""
+        self._coordinator = coordinator
+        self._client = client
+        self._unique_id = f"{entry.unique_id}_Response"
+        self._entry_id = entry.entry_id
+
+        self._state = None
+        self._state_attributes = {}
+        self._state_icon = None
+
+    @property
+    def name(self) -> str:
+        """Return the name of the switch."""
+        return "Incident Response"
+
+    @property
+    def icon(self) -> str:
+        """Return the icon to use in the frontend."""
+        if self._state_icon == "acknowledged":
+            return "mdi:run-fast"
+        if self._state_icon == "rejected":
+            return "mdi:account-off-outline"
+
+        return "mdi:forum"
+
+    @property
+    def is_on(self) -> bool:
+        """Get the assumed state of the switch."""
+        return self._state
+
+    @property
+    def unique_id(self) -> str:
+        """Return the unique ID for this switch."""
+        return self._unique_id
+
+    @property
+    def should_poll(self) -> bool:
+        """No polling needed."""
+        return False
+
+    @property
+    def available(self):
+        """Return if switch is available."""
+        return self._client.on_duty
+
+    @property
+    def device_state_attributes(self) -> object:
+        """Return available attributes for switch."""
+        attr = {}
+        if not self._state_attributes:
+            return attr
+
+        data = self._state_attributes
+        attr = {
+            key: data[key]
+            for key in (
+                "user_name",
+                "assigned_skill_ids",
+                "responded_at",
+                "start_time",
+                "status",
+                "reported_status",
+                "arrived_at_station",
+                "available_at_incident_creation",
+                "active_duty_function_ids",
+            )
+            if key in data
+        }
+
+        _LOGGER.debug("Set attributes of entity 'Response Switch' to '%s'", attr)
+        return attr
+
+    async def async_turn_on(self, **kwargs) -> None:
+        """Send Acknowlegde response status."""
+        await self.async_set_response(True)
+
+    async def async_turn_off(self, **kwargs) -> None:
+        """Send Reject response status."""
+        await self.async_set_response(False)
+
+    async def async_set_response(self, value) -> None:
+        """Send response status."""
+        if not self._client.on_duty:
+            _LOGGER.debug(
+                "Cannot send incident response when not on duty",
+            )
+            return
+
+        await self._client.async_set_response(value)
+        self.client_update()
+
+    async def async_added_to_hass(self) -> None:
+        """Register update callback."""
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                f"{FIRESERVICEROTA_DOMAIN}_{self._entry_id}_update",
+                self.client_update,
+            )
+        )
+        self.async_on_remove(self._coordinator.async_add_listener(self.on_duty_update))
+
+    @callback
+    def on_duty_update(self):
+        """Trigger on a on duty update."""
+        self.async_write_ha_state()
+
+    @callback
+    def client_update(self) -> None:
+        """Handle updated incident data from the client."""
+        self.async_schedule_update_ha_state(True)
+
+    async def async_update(self) -> bool:
+        """Update FireServiceRota response data."""
+        data = await self._client.async_response_update()
+
+        if not data or "status" not in data:
+            return
+
+        self._state = data["status"] == "acknowledged"
+        self._state_attributes = data
+        self._state_icon = data["status"]
+
+        _LOGGER.debug("Set state of entity 'Response Switch' to '%s'", self._state)