From 514e826fed5bd601223022920a3bde38a09b8b6c Mon Sep 17 00:00:00 2001 From: Jelte Zeilstra Date: Sat, 16 Jul 2022 20:39:11 +0200 Subject: [PATCH] Add install UniFi device update feature (#75302) * Add install UniFi device update feature * Add tests for install UniFi device update feature * Fix type error * Process review feedback * Process review feedback --- homeassistant/components/unifi/update.py | 13 +- tests/components/unifi/test_update.py | 165 +++++++++++++++-------- 2 files changed, 118 insertions(+), 60 deletions(-) diff --git a/homeassistant/components/unifi/update.py b/homeassistant/components/unifi/update.py index 09720f15f84..24967e043d9 100644 --- a/homeassistant/components/unifi/update.py +++ b/homeassistant/components/unifi/update.py @@ -2,6 +2,7 @@ from __future__ import annotations import logging +from typing import Any from homeassistant.components.update import ( DOMAIN, @@ -71,7 +72,6 @@ class UniFiDeviceUpdateEntity(UniFiBase, UpdateEntity): 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.""" @@ -79,6 +79,11 @@ class UniFiDeviceUpdateEntity(UniFiBase, UpdateEntity): self.device = self._item + self._attr_supported_features = UpdateEntityFeature.PROGRESS + + if self.controller.site_role == "admin": + self._attr_supported_features |= UpdateEntityFeature.INSTALL + @property def name(self) -> str: """Return the name of the device.""" @@ -126,3 +131,9 @@ class UniFiDeviceUpdateEntity(UniFiBase, UpdateEntity): async def options_updated(self) -> None: """No action needed.""" + + async def async_install( + self, version: str | None, backup: bool, **kwargs: Any + ) -> None: + """Install an update.""" + await self.controller.api.devices.upgrade(self.device.mac) diff --git a/tests/components/unifi/test_update.py b/tests/components/unifi/test_update.py index b5eb9c1d02e..7bbeb65497e 100644 --- a/tests/components/unifi/test_update.py +++ b/tests/components/unifi/test_update.py @@ -1,23 +1,59 @@ """The tests for the UniFi Network update platform.""" +from copy import deepcopy from aiounifi.controller import MESSAGE_DEVICE from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING +from yarl import URL +from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN from homeassistant.components.update import ( ATTR_IN_PROGRESS, ATTR_INSTALLED_VERSION, ATTR_LATEST_VERSION, DOMAIN as UPDATE_DOMAIN, + SERVICE_INSTALL, UpdateDeviceClass, + UpdateEntityFeature, ) from homeassistant.const import ( ATTR_DEVICE_CLASS, + ATTR_ENTITY_ID, + ATTR_SUPPORTED_FEATURES, STATE_OFF, STATE_ON, STATE_UNAVAILABLE, ) -from .test_controller import setup_unifi_integration +from .test_controller import DESCRIPTION, setup_unifi_integration + +DEVICE_1 = { + "board_rev": 3, + "device_id": "mock-id", + "ip": "10.0.1.1", + "last_seen": 1562600145, + "mac": "00:00:00:00:01:01", + "model": "US16P150", + "name": "Device 1", + "next_interval": 20, + "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", + "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", +} async def test_no_entities(hass, aioclient_mock): @@ -31,41 +67,11 @@ 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", - } + device_1 = deepcopy(DEVICE_1) await setup_unifi_integration( hass, aioclient_mock, - devices_response=[device_1, device_2], + devices_response=[device_1, DEVICE_2], ) assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 2 @@ -76,6 +82,10 @@ async def test_device_updates( 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 + assert ( + device_1_state.attributes[ATTR_SUPPORTED_FEATURES] + == UpdateEntityFeature.PROGRESS | UpdateEntityFeature.INSTALL + ) device_2_state = hass.states.get("update.device_2") assert device_2_state.state == STATE_OFF @@ -83,6 +93,10 @@ async def test_device_updates( 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 + assert ( + device_2_state.attributes[ATTR_SUPPORTED_FEATURES] + == UpdateEntityFeature.PROGRESS | UpdateEntityFeature.INSTALL + ) # Simulate start of update @@ -122,46 +136,79 @@ async def test_device_updates( 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", - } +async def test_not_admin(hass, aioclient_mock): + """Test that the INSTALL feature is not available on a non-admin account.""" + description = deepcopy(DESCRIPTION) + description[0]["site_role"] = "not admin" await setup_unifi_integration( hass, aioclient_mock, - devices_response=[device], + site_description=description, + devices_response=[DEVICE_1], ) assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1 - assert hass.states.get("update.device").state == STATE_ON + device_state = hass.states.get("update.device_1") + assert device_state.state == STATE_ON + assert ( + device_state.attributes[ATTR_SUPPORTED_FEATURES] == UpdateEntityFeature.PROGRESS + ) + + +async def test_install(hass, aioclient_mock): + """Test the device update install call.""" + config_entry = await setup_unifi_integration( + hass, aioclient_mock, devices_response=[DEVICE_1] + ) + + assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1 + device_state = hass.states.get("update.device_1") + assert device_state.state == STATE_ON + + controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] + url = f"https://{controller.host}:1234/api/s/{controller.site}/cmd/devmgr" + aioclient_mock.clear_requests() + aioclient_mock.post(url) + + await hass.services.async_call( + UPDATE_DOMAIN, + SERVICE_INSTALL, + {ATTR_ENTITY_ID: "update.device_1"}, + blocking=True, + ) + await hass.async_block_till_done() + + assert aioclient_mock.call_count == 1 + assert aioclient_mock.mock_calls[0] == ( + "post", + URL(url), + {"cmd": "upgrade", "mac": "00:00:00:00:01:01"}, + {}, + ) + + +async def test_controller_state_change( + hass, aioclient_mock, mock_unifi_websocket, mock_device_registry +): + """Verify entities state reflect on controller becoming unavailable.""" + await setup_unifi_integration( + hass, + aioclient_mock, + devices_response=[DEVICE_1], + ) + + assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1 + assert hass.states.get("update.device_1").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 + assert hass.states.get("update.device_1").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 + assert hass.states.get("update.device_1").state == STATE_ON