UniFi device tracker restore clients (#25532)

This commit is contained in:
Robert Svensson 2019-07-29 23:13:04 +02:00 committed by GitHub
parent 3a91c8f285
commit 891f19b43f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 53 additions and 3 deletions

View File

@ -6,12 +6,13 @@ import voluptuous as vol
from homeassistant import config_entries
from homeassistant.components import unifi
from homeassistant.components.device_tracker import PLATFORM_SCHEMA
from homeassistant.components.device_tracker import DOMAIN, PLATFORM_SCHEMA
from homeassistant.components.device_tracker.config_entry import ScannerEntity
from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER
from homeassistant.core import callback
from homeassistant.const import (
CONF_HOST, CONF_USERNAME, CONF_PASSWORD, CONF_PORT, CONF_VERIFY_SSL)
from homeassistant.helpers import entity_registry
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.dispatcher import async_dispatcher_connect
@ -80,6 +81,23 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
controller = hass.data[unifi.DOMAIN][controller_id]
tracked = {}
registry = await entity_registry.async_get_registry(hass)
# Restore clients that is not a part of active clients list.
for entity in registry.entities.values():
if entity.config_entry_id == config_entry.entry_id and \
entity.domain == DOMAIN:
mac, _ = entity.unique_id.split('-', 1)
if mac in controller.api.clients or \
mac not in controller.api.clients_all:
continue
client = controller.api.clients_all[mac]
controller.api.clients.process_raw([client.raw])
@callback
def update_controller():
"""Update the values of the controller."""

View File

@ -6,7 +6,7 @@ from datetime import timedelta
import pytest
from aiounifi.clients import Clients
from aiounifi.clients import Clients, ClientsAll
from aiounifi.devices import Devices
from homeassistant import config_entries
@ -15,6 +15,7 @@ from homeassistant.components.unifi.const import (
CONF_CONTROLLER, CONF_SITE_ID, UNIFI_CONFIG)
from homeassistant.const import (
CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL)
from homeassistant.helpers import entity_registry
from homeassistant.setup import async_setup_component
import homeassistant.components.device_tracker as device_tracker
@ -75,6 +76,7 @@ def mock_controller(hass):
controller.mock_client_responses = deque()
controller.mock_device_responses = deque()
controller.mock_client_all_responses = deque()
async def mock_request(method, path, **kwargs):
kwargs['method'] = method
@ -84,10 +86,13 @@ def mock_controller(hass):
return controller.mock_client_responses.popleft()
if path == 's/{site}/stat/device':
return controller.mock_device_responses.popleft()
if path == 's/{site}/rest/user':
return controller.mock_client_all_responses.popleft()
return None
controller.api.clients = Clients({}, mock_request)
controller.api.devices = Devices({}, mock_request)
controller.api.clients_all = ClientsAll({}, mock_request)
return controller
@ -98,7 +103,7 @@ async def setup_controller(hass, mock_controller):
hass.data[unifi.DOMAIN] = {CONTROLLER_ID: mock_controller}
config_entry = config_entries.ConfigEntry(
1, unifi.DOMAIN, 'Mock Title', ENTRY_CONFIG, 'test',
config_entries.CONN_CLASS_LOCAL_POLL)
config_entries.CONN_CLASS_LOCAL_POLL, entry_id=1)
mock_controller.config_entry = config_entry
await mock_controller.async_update()
@ -159,3 +164,30 @@ async def test_tracked_devices(hass, mock_controller):
device_1 = hass.states.get('device_tracker.client_1')
assert device_1.state == 'home'
async def test_restoring_client(hass, mock_controller):
"""Test the update_items function with some clients."""
mock_controller.mock_client_responses.append([CLIENT_2])
mock_controller.mock_device_responses.append({})
mock_controller.mock_client_all_responses.append([CLIENT_1])
mock_controller.unifi_config = {
unifi.CONF_BLOCK_CLIENT: True
}
registry = await entity_registry.async_get_registry(hass)
registry.async_get_or_create(
device_tracker.DOMAIN, unifi_dt.UNIFI_DOMAIN,
'{}-mock-site'.format(CLIENT_1['mac']),
suggested_object_id=CLIENT_1['hostname'], config_entry_id=1)
registry.async_get_or_create(
device_tracker.DOMAIN, unifi_dt.UNIFI_DOMAIN,
'{}-mock-site'.format(CLIENT_2['mac']),
suggested_object_id=CLIENT_2['hostname'], config_entry_id=1)
await setup_controller(hass, mock_controller)
assert len(mock_controller.mock_requests) == 3
assert len(hass.states.async_all()) == 4
device_1 = hass.states.get('device_tracker.client_1')
assert device_1 is not None