Add update entity to RainMachine (#76100)

* Add update entity to RainMachine

* Fix tests

* Cleanup

* Test missing controller diagnostics

* Code review

.coveragerc

@@ -980,6 +980,7 @@ omit =
     homeassistant/components/rainmachine/model.py
     homeassistant/components/rainmachine/sensor.py
     homeassistant/components/rainmachine/switch.py
+    homeassistant/components/rainmachine/update.py
     homeassistant/components/rainmachine/util.py
     homeassistant/components/raspyrfm/*
     homeassistant/components/recollect_waste/__init__.py

homeassistant/components/rainmachine/__init__.py

@@ -36,6 +36,8 @@ from homeassistant.util.network import is_ip_address
 from .config_flow import get_client_controller
 from .const import (
     CONF_ZONE_RUN_TIME,
+    DATA_API_VERSIONS,
+    DATA_MACHINE_FIRMWARE_UPDATE_STATUS,
     DATA_PROGRAMS,
     DATA_PROVISION_SETTINGS,
     DATA_RESTRICTIONS_CURRENT,
@@ -51,7 +53,13 @@ DEFAULT_SSL = True
 
 CONFIG_SCHEMA = cv.removed(DOMAIN, raise_if_present=False)
 
-PLATFORMS = [Platform.BINARY_SENSOR, Platform.BUTTON, Platform.SENSOR, Platform.SWITCH]
+PLATFORMS = [
+    Platform.BINARY_SENSOR,
+    Platform.BUTTON,
+    Platform.SENSOR,
+    Platform.SWITCH,
+    Platform.UPDATE,
+]
 
 CONF_CONDITION = "condition"
 CONF_DEWPOINT = "dewpoint"
@@ -124,8 +132,10 @@ SERVICE_RESTRICT_WATERING_SCHEMA = SERVICE_SCHEMA.extend(
 )
 
 COORDINATOR_UPDATE_INTERVAL_MAP = {
-    DATA_PROVISION_SETTINGS: timedelta(minutes=1),
+    DATA_API_VERSIONS: timedelta(minutes=1),
+    DATA_MACHINE_FIRMWARE_UPDATE_STATUS: timedelta(seconds=15),
     DATA_PROGRAMS: timedelta(seconds=30),
+    DATA_PROVISION_SETTINGS: timedelta(minutes=1),
     DATA_RESTRICTIONS_CURRENT: timedelta(minutes=1),
     DATA_RESTRICTIONS_UNIVERSAL: timedelta(minutes=1),
     DATA_ZONES: timedelta(seconds=15),
@@ -215,7 +225,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         data: dict = {}
 
         try:
-            if api_category == DATA_PROGRAMS:
+            if api_category == DATA_API_VERSIONS:
+                data = await controller.api.versions()
+            elif api_category == DATA_MACHINE_FIRMWARE_UPDATE_STATUS:
+                data = await controller.machine.get_firmware_update_status()
+            elif api_category == DATA_PROGRAMS:
                 data = await controller.programs.all(include_inactive=True)
             elif api_category == DATA_PROVISION_SETTINGS:
                 data = await controller.provisioning.settings()
@@ -414,23 +428,32 @@ class RainMachineEntity(CoordinatorEntity):
         """Initialize."""
         super().__init__(data.coordinators[description.api_category])
 
-        self._attr_device_info = DeviceInfo(
-            identifiers={(DOMAIN, data.controller.mac)},
-            configuration_url=f"https://{entry.data[CONF_IP_ADDRESS]}:{entry.data[CONF_PORT]}",
-            connections={(dr.CONNECTION_NETWORK_MAC, data.controller.mac)},
-            name=str(data.controller.name).capitalize(),
-            manufacturer="RainMachine",
-            model=(
-                f"Version {data.controller.hardware_version} "
-                f"(API: {data.controller.api_version})"
-            ),
-            sw_version=data.controller.software_version,
-        )
         self._attr_extra_state_attributes = {}
         self._attr_unique_id = f"{data.controller.mac}_{description.key}"
+        self._entry = entry
         self._data = data
+        self._version_coordinator = data.coordinators[DATA_API_VERSIONS]
         self.entity_description = description
 
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information about this controller."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._data.controller.mac)},
+            configuration_url=(
+                f"https://{self._entry.data[CONF_IP_ADDRESS]}:"
+                f"{self._entry.data[CONF_PORT]}"
+            ),
+            connections={(dr.CONNECTION_NETWORK_MAC, self._data.controller.mac)},
+            name=str(self._data.controller.name).capitalize(),
+            manufacturer="RainMachine",
+            model=(
+                f"Version {self._version_coordinator.data['hwVer']} "
+                f"(API: {self._version_coordinator.data['apiVer']})"
+            ),
+            sw_version=self._version_coordinator.data["swVer"],
+        )
+
     @callback
     def _handle_coordinator_update(self) -> None:
         """Respond to a DataUpdateCoordinator update."""
@@ -440,6 +463,11 @@ class RainMachineEntity(CoordinatorEntity):
     async def async_added_to_hass(self) -> None:
         """When entity is added to hass."""
         await super().async_added_to_hass()
+        self.async_on_remove(
+            self._version_coordinator.async_add_listener(
+                self._handle_coordinator_update, self.coordinator_context
+            )
+        )
         self.update_from_latest_data()
 
     @callback

homeassistant/components/rainmachine/const.py

@@ -7,6 +7,8 @@ DOMAIN = "rainmachine"
 
 CONF_ZONE_RUN_TIME = "zone_run_time"
 
+DATA_API_VERSIONS = "api.versions"
+DATA_MACHINE_FIRMWARE_UPDATE_STATUS = "machine.firmware_update_status"
 DATA_PROGRAMS = "programs"
 DATA_PROVISION_SETTINGS = "provision.settings"
 DATA_RESTRICTIONS_CURRENT = "restrictions.current"

homeassistant/components/rainmachine/diagnostics.py

 """Diagnostics support for RainMachine."""
 from __future__ import annotations
 
-import asyncio
 from typing import Any
 
-from regenmaschine.errors import RequestError
+from regenmaschine.errors import RainMachineError
 
 from homeassistant.components.diagnostics import async_redact_data
 from homeassistant.config_entries import ConfigEntry
@@ -17,7 +16,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 
 from . import RainMachineData
-from .const import DOMAIN
+from .const import DOMAIN, LOGGER
 
 CONF_STATION_ID = "stationID"
 CONF_STATION_NAME = "stationName"
@@ -42,13 +41,11 @@ async def async_get_config_entry_diagnostics(
     """Return diagnostics for a config entry."""
     data: RainMachineData = hass.data[DOMAIN][entry.entry_id]
 
-    controller_tasks = {
-        "versions": data.controller.api.versions(),
-        "current_diagnostics": data.controller.diagnostics.current(),
-    }
-    controller_results = await asyncio.gather(
-        *controller_tasks.values(), return_exceptions=True
-    )
+    try:
+        controller_diagnostics = await data.controller.diagnostics.current()
+    except RainMachineError:
+        LOGGER.warning("Unable to download controller-specific diagnostics")
+        controller_diagnostics = None
 
     return {
         "entry": {
@@ -64,10 +61,6 @@ async def async_get_config_entry_diagnostics(
                 },
                 TO_REDACT,
             ),
-            "controller": {
-                category: result
-                for category, result in zip(controller_tasks, controller_results)
-                if not isinstance(result, RequestError)
-            },
+            "controller_diagnostics": controller_diagnostics,
         },
     }

homeassistant/components/rainmachine/update.py

+"""Support for RainMachine updates."""
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any
+
+from regenmaschine.errors import RequestError
+
+from homeassistant.components.update import (
+    UpdateDeviceClass,
+    UpdateEntity,
+    UpdateEntityFeature,
+)
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from . import RainMachineData, RainMachineEntity
+from .const import DATA_MACHINE_FIRMWARE_UPDATE_STATUS, DOMAIN
+from .model import RainMachineEntityDescription
+
+
+class UpdateStates(Enum):
+    """Define an enum for update states."""
+
+    IDLE = 1
+    CHECKING = 2
+    DOWNLOADING = 3
+    UPGRADING = 4
+    ERROR = 5
+    REBOOT = 6
+
+
+UPDATE_STATE_MAP = {
+    1: UpdateStates.IDLE,
+    2: UpdateStates.CHECKING,
+    3: UpdateStates.DOWNLOADING,
+    4: UpdateStates.UPGRADING,
+    5: UpdateStates.ERROR,
+    6: UpdateStates.REBOOT,
+}
+
+
+UPDATE_DESCRIPTION = RainMachineEntityDescription(
+    key="update",
+    name="Firmware",
+    api_category=DATA_MACHINE_FIRMWARE_UPDATE_STATUS,
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
+    """Set up WLED update based on a config entry."""
+    data: RainMachineData = hass.data[DOMAIN][entry.entry_id]
+
+    async_add_entities([RainMachineUpdateEntity(entry, data, UPDATE_DESCRIPTION)])
+
+
+class RainMachineUpdateEntity(RainMachineEntity, UpdateEntity):
+    """Define a RainMachine update entity."""
+
+    _attr_device_class = UpdateDeviceClass.FIRMWARE
+    _attr_supported_features = (
+        UpdateEntityFeature.INSTALL
+        | UpdateEntityFeature.PROGRESS
+        | UpdateEntityFeature.SPECIFIC_VERSION
+    )
+
+    async def async_install(
+        self, version: str | None, backup: bool, **kwargs: Any
+    ) -> None:
+        """Install an update."""
+        try:
+            await self._data.controller.machine.update_firmware()
+        except RequestError as err:
+            raise HomeAssistantError(f"Error while updating firmware: {err}") from err
+
+        await self.coordinator.async_refresh()
+
+    @callback
+    def update_from_latest_data(self) -> None:
+        """Update the state."""
+        if version := self._version_coordinator.data["swVer"]:
+            self._attr_installed_version = version
+        else:
+            self._attr_installed_version = None
+
+        data = self.coordinator.data
+
+        if not data["update"]:
+            self._attr_in_progress = False
+            self._attr_latest_version = self._attr_installed_version
+            return
+
+        self._attr_in_progress = UPDATE_STATE_MAP[data["updateStatus"]] in (
+            UpdateStates.DOWNLOADING,
+            UpdateStates.UPGRADING,
+            UpdateStates.REBOOT,
+        )
+        self._attr_latest_version = data["packageDetails"]["newVersion"]

tests/components/rainmachine/conftest.py

@@ -41,6 +41,7 @@ def controller_fixture(
     controller_mac,
     data_api_versions,
     data_diagnostics_current,
+    data_machine_firmare_update_status,
     data_programs,
     data_provision_settings,
     data_restrictions_current,
@@ -59,6 +60,9 @@ def controller_fixture(
 
     controller.api.versions.return_value = data_api_versions
     controller.diagnostics.current.return_value = data_diagnostics_current
+    controller.machine.get_firmware_update_status.return_value = (
+        data_machine_firmare_update_status
+    )
     controller.programs.all.return_value = data_programs
     controller.provisioning.settings.return_value = data_provision_settings
     controller.restrictions.current.return_value = data_restrictions_current
@@ -86,6 +90,14 @@ def data_diagnostics_current_fixture():
     return json.loads(load_fixture("diagnostics_current_data.json", "rainmachine"))
 
 
+@pytest.fixture(name="data_machine_firmare_update_status", scope="session")
+def data_machine_firmare_update_status_fixture():
+    """Define machine firmware update status data."""
+    return json.loads(
+        load_fixture("machine_firmware_update_status_data.json", "rainmachine")
+    )
+
+
 @pytest.fixture(name="data_programs", scope="session")
 def data_programs_fixture():
     """Define program data."""

tests/components/rainmachine/fixtures/machine_firmware_update_status_data.json

+{
+  "lastUpdateCheckTimestamp": 1657825288,
+  "packageDetails": [],
+  "update": false,
+  "lastUpdateCheck": "2022-07-14 13:01:28",
+  "updateStatus": 1
+}