Create Update entities for Ubiquiti network devices (#71700)

Co-authored-by: Robert Svensson <Kane610@users.noreply.github.com>

homeassistant/components/unifi/controller.py

@@ -74,7 +74,7 @@ from .switch import BLOCK_SWITCH, POE_SWITCH
 
 RETRY_TIMER = 15
 CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1)
-PLATFORMS = [Platform.DEVICE_TRACKER, Platform.SENSOR, Platform.SWITCH]
+PLATFORMS = [Platform.DEVICE_TRACKER, Platform.SENSOR, Platform.SWITCH, Platform.UPDATE]
 
 CLIENT_CONNECTED = (
     WIRED_CLIENT_CONNECTED,

homeassistant/components/unifi/device_tracker.py

@@ -429,6 +429,16 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
 
         return attributes
 
+    @property
+    def ip_address(self) -> str:
+        """Return the primary ip address of the device."""
+        return self.device.ip
+
+    @property
+    def mac_address(self) -> str:
+        """Return the mac address of the device."""
+        return self.device.mac
+
     async def options_updated(self) -> None:
         """Config entry options are updated, remove entity if option is disabled."""
         if not self.controller.option_track_devices:

homeassistant/components/unifi/update.py

@@ -0,0 +1,128 @@
+"""Update entities for Ubiquiti network devices."""
+from __future__ import annotations
+
+import logging
+
+from homeassistant.components.update import (
+    DOMAIN,
+    UpdateDeviceClass,
+    UpdateEntity,
+    UpdateEntityFeature,
+)
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import ATTR_NAME
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN
+from .unifi_entity_base import UniFiBase
+
+LOGGER = logging.getLogger(__name__)
+
+DEVICE_UPDATE = "device_update"
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up update entities for UniFi Network integration."""
+    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+    controller.entities[DOMAIN] = {DEVICE_UPDATE: set()}
+
+    @callback
+    def items_added(
+        clients: set = controller.api.clients, devices: set = controller.api.devices
+    ) -> None:
+        """Add device update entities."""
+        add_device_update_entities(controller, async_add_entities, devices)
+
+    for signal in (controller.signal_update, controller.signal_options_update):
+        config_entry.async_on_unload(
+            async_dispatcher_connect(hass, signal, items_added)
+        )
+
+    items_added()
+
+
+@callback
+def add_device_update_entities(controller, async_add_entities, devices):
+    """Add new device update entities from the controller."""
+    entities = []
+
+    for mac in devices:
+        if mac in controller.entities[DOMAIN][UniFiDeviceUpdateEntity.TYPE]:
+            continue
+
+        device = controller.api.devices[mac]
+        entities.append(UniFiDeviceUpdateEntity(device, controller))
+
+    if entities:
+        async_add_entities(entities)
+
+
+class UniFiDeviceUpdateEntity(UniFiBase, UpdateEntity):
+    """Update entity for a UniFi network infrastructure device."""
+
+    DOMAIN = DOMAIN
+    TYPE = DEVICE_UPDATE
+    _attr_device_class = UpdateDeviceClass.FIRMWARE
+    _attr_supported_features = UpdateEntityFeature.PROGRESS
+
+    def __init__(self, device, controller):
+        """Set up device update entity."""
+        super().__init__(device, controller)
+
+        self.device = self._item
+
+    @property
+    def name(self) -> str:
+        """Return the name of the device."""
+        return self.device.name or self.device.model
+
+    @property
+    def unique_id(self) -> str:
+        """Return a unique identifier for this device."""
+        return f"{self.TYPE}-{self.device.mac}"
+
+    @property
+    def available(self) -> bool:
+        """Return if controller is available."""
+        return not self.device.disabled and self.controller.available
+
+    @property
+    def in_progress(self) -> bool:
+        """Update installation in progress."""
+        return self.device.state == 4
+
+    @property
+    def installed_version(self) -> str | None:
+        """Version currently in use."""
+        return self.device.version
+
+    @property
+    def latest_version(self) -> str | None:
+        """Latest version available for install."""
+        return self.device.upgrade_to_firmware or self.device.version
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return a device description for device registry."""
+        info = DeviceInfo(
+            connections={(CONNECTION_NETWORK_MAC, self.device.mac)},
+            manufacturer=ATTR_MANUFACTURER,
+            model=self.device.model,
+            sw_version=self.device.version,
+        )
+
+        if self.device.name:
+            info[ATTR_NAME] = self.device.name
+
+        return info
+
+    async def options_updated(self) -> None:
+        """No action needed."""

tests/components/unifi/test_device_tracker.py

@@ -499,6 +499,7 @@ async def test_option_track_devices(hass, aioclient_mock, mock_device_registry):
         "board_rev": 3,
         "device_id": "mock-id",
         "last_seen": 1562600145,
+        "ip": "10.0.1.1",
         "mac": "00:00:00:00:01:01",
         "model": "US16P150",
         "name": "Device",

tests/components/unifi/test_diagnostics.py

@@ -14,6 +14,7 @@ from homeassistant.components.unifi.switch import (
     OUTLET_SWITCH,
     POE_SWITCH,
 )
+from homeassistant.components.unifi.update import DEVICE_UPDATE
 from homeassistant.const import Platform
 
 from .test_controller import setup_unifi_integration
@@ -161,6 +162,9 @@ async def test_entry_diagnostics(hass, hass_client, aioclient_mock):
                 POE_SWITCH: ["00:00:00:00:00:00"],
                 OUTLET_SWITCH: [],
             },
+            str(Platform.UPDATE): {
+                DEVICE_UPDATE: ["00:00:00:00:00:01"],
+            },
         },
         "clients": {
             "00:00:00:00:00:00": {

tests/components/unifi/test_update.py

@@ -0,0 +1,167 @@
+"""The tests for the UniFi Network update platform."""
+
+from aiounifi.controller import MESSAGE_DEVICE
+from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
+
+from homeassistant.components.update import (
+    ATTR_IN_PROGRESS,
+    ATTR_INSTALLED_VERSION,
+    ATTR_LATEST_VERSION,
+    DOMAIN as UPDATE_DOMAIN,
+    UpdateDeviceClass,
+)
+from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
+    STATE_OFF,
+    STATE_ON,
+    STATE_UNAVAILABLE,
+)
+
+from .test_controller import setup_unifi_integration
+
+
+async def test_no_entities(hass, aioclient_mock):
+    """Test the update_clients function when no clients are found."""
+    await setup_unifi_integration(hass, aioclient_mock)
+
+    assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 0
+
+
+async def test_device_updates(
+    hass, aioclient_mock, mock_unifi_websocket, mock_device_registry
+):
+    """Test the update_items function with some devices."""
+    device_1 = {
+        "board_rev": 3,
+        "device_id": "mock-id",
+        "has_fan": True,
+        "fan_level": 0,
+        "ip": "10.0.1.1",
+        "last_seen": 1562600145,
+        "mac": "00:00:00:00:01:01",
+        "model": "US16P150",
+        "name": "Device 1",
+        "next_interval": 20,
+        "overheating": True,
+        "state": 1,
+        "type": "usw",
+        "upgradable": True,
+        "version": "4.0.42.10433",
+        "upgrade_to_firmware": "4.3.17.11279",
+    }
+    device_2 = {
+        "board_rev": 3,
+        "device_id": "mock-id",
+        "has_fan": True,
+        "ip": "10.0.1.2",
+        "mac": "00:00:00:00:01:02",
+        "model": "US16P150",
+        "name": "Device 2",
+        "next_interval": 20,
+        "state": 0,
+        "type": "usw",
+        "version": "4.0.42.10433",
+    }
+    await setup_unifi_integration(
+        hass,
+        aioclient_mock,
+        devices_response=[device_1, device_2],
+    )
+
+    assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 2
+
+    device_1_state = hass.states.get("update.device_1")
+    assert device_1_state.state == STATE_ON
+    assert device_1_state.attributes[ATTR_INSTALLED_VERSION] == "4.0.42.10433"
+    assert device_1_state.attributes[ATTR_LATEST_VERSION] == "4.3.17.11279"
+    assert device_1_state.attributes[ATTR_IN_PROGRESS] is False
+    assert device_1_state.attributes[ATTR_DEVICE_CLASS] == UpdateDeviceClass.FIRMWARE
+
+    device_2_state = hass.states.get("update.device_2")
+    assert device_2_state.state == STATE_OFF
+    assert device_2_state.attributes[ATTR_INSTALLED_VERSION] == "4.0.42.10433"
+    assert device_2_state.attributes[ATTR_LATEST_VERSION] == "4.0.42.10433"
+    assert device_2_state.attributes[ATTR_IN_PROGRESS] is False
+    assert device_2_state.attributes[ATTR_DEVICE_CLASS] == UpdateDeviceClass.FIRMWARE
+
+    # Simulate start of update
+
+    device_1["state"] = 4
+    mock_unifi_websocket(
+        data={
+            "meta": {"message": MESSAGE_DEVICE},
+            "data": [device_1],
+        }
+    )
+    await hass.async_block_till_done()
+
+    device_1_state = hass.states.get("update.device_1")
+    assert device_1_state.state == STATE_ON
+    assert device_1_state.attributes[ATTR_INSTALLED_VERSION] == "4.0.42.10433"
+    assert device_1_state.attributes[ATTR_LATEST_VERSION] == "4.3.17.11279"
+    assert device_1_state.attributes[ATTR_IN_PROGRESS] is True
+
+    # Simulate update finished
+
+    device_1["state"] = "0"
+    device_1["version"] = "4.3.17.11279"
+    device_1["upgradable"] = False
+    del device_1["upgrade_to_firmware"]
+    mock_unifi_websocket(
+        data={
+            "meta": {"message": MESSAGE_DEVICE},
+            "data": [device_1],
+        }
+    )
+    await hass.async_block_till_done()
+
+    device_1_state = hass.states.get("update.device_1")
+    assert device_1_state.state == STATE_OFF
+    assert device_1_state.attributes[ATTR_INSTALLED_VERSION] == "4.3.17.11279"
+    assert device_1_state.attributes[ATTR_LATEST_VERSION] == "4.3.17.11279"
+    assert device_1_state.attributes[ATTR_IN_PROGRESS] is False
+
+
+async def test_controller_state_change(
+    hass, aioclient_mock, mock_unifi_websocket, mock_device_registry
+):
+    """Verify entities state reflect on controller becoming unavailable."""
+    device = {
+        "board_rev": 3,
+        "device_id": "mock-id",
+        "has_fan": True,
+        "fan_level": 0,
+        "ip": "10.0.1.1",
+        "last_seen": 1562600145,
+        "mac": "00:00:00:00:01:01",
+        "model": "US16P150",
+        "name": "Device",
+        "next_interval": 20,
+        "overheating": True,
+        "state": 1,
+        "type": "usw",
+        "upgradable": True,
+        "version": "4.0.42.10433",
+        "upgrade_to_firmware": "4.3.17.11279",
+    }
+
+    await setup_unifi_integration(
+        hass,
+        aioclient_mock,
+        devices_response=[device],
+    )
+
+    assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1
+    assert hass.states.get("update.device").state == STATE_ON
+
+    # Controller unavailable
+    mock_unifi_websocket(state=STATE_DISCONNECTED)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("update.device").state == STATE_UNAVAILABLE
+
+    # Controller available
+    mock_unifi_websocket(state=STATE_RUNNING)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("update.device").state == STATE_ON