UniFi - Update device registry after device has been upgraded (#35430)

* Update device registry if received event saying device has been updated

* Make sure event was the last updated
This commit is contained in:
Robert Svensson 2020-05-10 17:14:45 +02:00 committed by GitHub
parent 4395fe9298
commit c71b6c8a71
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 49 additions and 0 deletions

View File

@ -4,6 +4,9 @@ import logging
from aiounifi.api import SOURCE_DATA, SOURCE_EVENT
from aiounifi.events import (
ACCESS_POINT_UPGRADED,
GATEWAY_UPGRADED,
SWITCH_UPGRADED,
WIRED_CLIENT_CONNECTED,
WIRELESS_CLIENT_CONNECTED,
WIRELESS_CLIENT_ROAM,
@ -50,6 +53,8 @@ CLIENT_STATIC_ATTRIBUTES = [
"oui",
]
DEVICE_UPGRADED = (ACCESS_POINT_UPGRADED, GATEWAY_UPGRADED, SWITCH_UPGRADED)
WIRED_CONNECTION = (WIRED_CLIENT_CONNECTED,)
WIRELESS_CONNECTION = (
WIRELESS_CLIENT_CONNECTED,
@ -299,6 +304,13 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
dt_util.utcnow() + timedelta(seconds=self.device.next_interval + 10),
)
elif (
self.device.last_updated == SOURCE_EVENT
and self.device.event.event in DEVICE_UPGRADED
):
self.hass.async_create_task(self.async_update_device_registry())
return
super().async_update_callback()
@property
@ -341,6 +353,14 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
return info
async def async_update_device_registry(self) -> None:
"""Update device registry."""
device_registry = await self.hass.helpers.device_registry.async_get_registry()
device_registry.async_get_or_create(
config_entry_id=self.controller.config_entry.entry_id, **self.device_info
)
@property
def device_state_attributes(self):
"""Return the device state attributes."""

View File

@ -136,6 +136,19 @@ EVENT_CLIENT_1_WIRELESS_DISCONNECTED = {
"_id": "5ea32ff730c49e00f90dca1a",
}
EVENT_DEVICE_2_UPGRADED = {
"_id": "5eae7fe02ab79c00f9d38960",
"datetime": "2020-05-09T20:06:37Z",
"key": "EVT_SW_Upgraded",
"msg": f'Switch[{DEVICE_2["mac"]}] was upgraded from "{DEVICE_2["version"]}" to "4.3.13.11253"',
"subsystem": "lan",
"sw": DEVICE_2["mac"],
"sw_name": DEVICE_2["name"],
"time": 1589054797635,
"version_from": {DEVICE_2["version"]},
"version_to": "4.3.13.11253",
}
async def test_platform_manually_configured(hass):
"""Test that nothing happens when configuring unifi through device tracker platform."""
@ -298,6 +311,22 @@ async def test_tracked_devices(hass):
device_1 = hass.states.get("device_tracker.device_1")
assert device_1.state == STATE_UNAVAILABLE
# Update device registry when device is upgraded
device_2_copy = copy(DEVICE_2)
device_2_copy["version"] = EVENT_DEVICE_2_UPGRADED["version_to"]
message = {"meta": {"message": MESSAGE_DEVICE}, "data": [device_2_copy]}
controller.api.message_handler(message)
event = {"meta": {"message": MESSAGE_EVENT}, "data": [EVENT_DEVICE_2_UPGRADED]}
controller.api.message_handler(event)
await hass.async_block_till_done()
# Verify device registry has been updated
entity_registry = await hass.helpers.entity_registry.async_get_registry()
entry = entity_registry.async_get("device_tracker.device_2")
device_registry = await hass.helpers.device_registry.async_get_registry()
device = device_registry.async_get(entry.device_id)
assert device.sw_version == EVENT_DEVICE_2_UPGRADED["version_to"]
async def test_remove_clients(hass):
"""Test the remove_items function with some clients."""