evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

UniFi - use entity registry disabled_by to control available entities (#26141)

Browse source 
* Move ignoring logic to entity registry enabled default

* Handle config to option import better

* Properly enable and disable entity registry entries on changes from config entry options

* Fix balloobs comments

* Fix some tests

* Fix tests

* Simplify updating disable on entities

* Simplify device tracker update function

* Local entity disabled replaced on rebase

* Only alter entities with changed options

* Proper tracking of changed options

* Back to straightforward updating of disabled
This commit is contained in:
Robert Svensson 2019-08-31 22:04:04 +02:00 • committed by GitHub
parent 922522b089
commit d9ef92f6d2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 172 additions and 105 deletions
Download patch file Download diff file Expand all files Collapse all files

6
homeassistant/components/unifi/__init__.py
View file

@ -11,6 +11,9 @@ from .const import (
CONF_BLOCK_CLIENT,
CONF_CONTROLLER,
CONF_DETECTION_TIME,
CONF_DONT_TRACK_CLIENTS,
CONF_DONT_TRACK_DEVICES,
CONF_DONT_TRACK_WIRED_CLIENTS,
CONF_SITE_ID,
CONF_SSID_FILTER,
CONTROLLER_ID,
@ -20,9 +23,6 @@ from .const import (
from .controller import UniFiController
CONF_CONTROLLERS = "controllers"
CONF_DONT_TRACK_CLIENTS = "dont_track_clients"
CONF_DONT_TRACK_DEVICES = "dont_track_devices"
CONF_DONT_TRACK_WIRED_CLIENTS = "dont_track_wired_clients"
CONTROLLER_SCHEMA = vol.Schema(
{

4
homeassistant/components/unifi/const.py
View file

@ -18,6 +18,10 @@ CONF_TRACK_DEVICES = "track_devices"
CONF_TRACK_WIRED_CLIENTS = "track_wired_clients"
CONF_SSID_FILTER = "ssid_filter"
CONF_DONT_TRACK_CLIENTS = "dont_track_clients"
CONF_DONT_TRACK_DEVICES = "dont_track_devices"
CONF_DONT_TRACK_WIRED_CLIENTS = "dont_track_wired_clients"
DEFAULT_BLOCK_CLIENTS = []
DEFAULT_TRACK_CLIENTS = True
DEFAULT_TRACK_DEVICES = True

100
homeassistant/components/unifi/controller.py
View file

@ -18,6 +18,9 @@ from .const import (
CONF_BLOCK_CLIENT,
CONF_CONTROLLER,
CONF_DETECTION_TIME,
CONF_DONT_TRACK_CLIENTS,
CONF_DONT_TRACK_DEVICES,
CONF_DONT_TRACK_WIRED_CLIENTS,
CONF_TRACK_CLIENTS,
CONF_TRACK_DEVICES,
CONF_TRACK_WIRED_CLIENTS,
@ -30,6 +33,7 @@ from .const import (
DEFAULT_TRACK_WIRED_CLIENTS,
DEFAULT_DETECTION_TIME,
DEFAULT_SSID_FILTER,
DOMAIN,
LOGGER,
UNIFI_CONFIG,
)
@ -49,7 +53,6 @@ class UniFiController:
self._site_name = None
self._site_role = None
self.unifi_config = {}
@property
def host(self):
@ -116,11 +119,14 @@ class UniFiController:
return None
@property
def event_update(self):
def signal_update(self):
"""Event specific per UniFi entry to signal new data."""
return "unifi-update-{}".format(
CONTROLLER_ID.format(host=self.host, site=self.site)
)
return f"unifi-update-{CONTROLLER_ID.format(host=self.host, site=self.site)}"
@property
def signal_options_update(self):
"""Event specific per UniFi entry to signal new options."""
return f"unifi-options-{CONTROLLER_ID.format(host=self.host, site=self.site)}"
async def request_update(self):
"""Request an update."""
@ -164,7 +170,7 @@ class UniFiController:
LOGGER.info("Reconnected to controller %s", self.host)
self.available = True
async_dispatcher_send(self.hass, self.event_update)
async_dispatcher_send(self.hass, self.signal_update)
async def async_setup(self):
"""Set up a UniFi controller."""
@ -191,37 +197,9 @@ class UniFiController:
LOGGER.error("Unknown error connecting with UniFi controller: %s", err)
return False
for unifi_config in hass.data[UNIFI_CONFIG]:
if (
self.host == unifi_config[CONF_HOST]
and self.site_name == unifi_config[CONF_SITE_ID]
):
self.unifi_config = unifi_config
break
self.import_configuration()
options = dict(self.config_entry.options)
if CONF_BLOCK_CLIENT in self.unifi_config:
options[CONF_BLOCK_CLIENT] = self.unifi_config[CONF_BLOCK_CLIENT]
if CONF_TRACK_CLIENTS in self.unifi_config:
options[CONF_TRACK_CLIENTS] = self.unifi_config[CONF_TRACK_CLIENTS]
if CONF_TRACK_DEVICES in self.unifi_config:
options[CONF_TRACK_DEVICES] = self.unifi_config[CONF_TRACK_DEVICES]
if CONF_TRACK_WIRED_CLIENTS in self.unifi_config:
options[CONF_TRACK_WIRED_CLIENTS] = self.unifi_config[
CONF_TRACK_WIRED_CLIENTS
]
if CONF_DETECTION_TIME in self.unifi_config:
options[CONF_DETECTION_TIME] = self.unifi_config[CONF_DETECTION_TIME]
if CONF_SSID_FILTER in self.unifi_config:
options[CONF_SSID_FILTER] = self.unifi_config[CONF_SSID_FILTER]
hass.config_entries.async_update_entry(self.config_entry, options=options)
self.config_entry.add_update_listener(self.async_options_updated)
for platform in ["device_tracker", "switch"]:
hass.async_create_task(
@ -232,6 +210,56 @@ class UniFiController:
return True
@staticmethod
async def async_options_updated(hass, entry):
"""Triggered by config entry options updates."""
controller_id = CONTROLLER_ID.format(
host=entry.data[CONF_CONTROLLER][CONF_HOST],
site=entry.data[CONF_CONTROLLER][CONF_SITE_ID],
)
controller = hass.data[DOMAIN][controller_id]
async_dispatcher_send(hass, controller.signal_options_update)
def import_configuration(self):
"""Import configuration to config entry options."""
unifi_config = {}
for config in self.hass.data[UNIFI_CONFIG]:
if (
self.host == config[CONF_HOST]
and self.site_name == config[CONF_SITE_ID]
):
unifi_config = config
break
old_options = dict(self.config_entry.options)
new_options = {}
for config, option in (
(CONF_BLOCK_CLIENT, CONF_BLOCK_CLIENT),
(CONF_DONT_TRACK_CLIENTS, CONF_TRACK_CLIENTS),
(CONF_DONT_TRACK_WIRED_CLIENTS, CONF_TRACK_WIRED_CLIENTS),
(CONF_DONT_TRACK_DEVICES, CONF_TRACK_DEVICES),
(CONF_DETECTION_TIME, CONF_DETECTION_TIME),
(CONF_SSID_FILTER, CONF_SSID_FILTER),
):
if config in unifi_config:
if config == option and unifi_config[
config
] != self.config_entry.options.get(option):
new_options[option] = unifi_config[config]
elif config != option and (
option not in self.config_entry.options
or unifi_config[config] == self.config_entry.options.get(option)
):
new_options[option] = not unifi_config[config]
if new_options:
options = {**old_options, **new_options}
self.hass.config_entries.async_update_entry(
self.config_entry, options=options
)
async def async_reset(self):
"""Reset this controller to default state.

125
homeassistant/components/unifi/device_tracker.py
View file

@ -20,6 +20,7 @@ from homeassistant.const import (
from homeassistant.helpers import entity_registry
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY
import homeassistant.helpers.config_validation as cv
import homeassistant.util.dt as dt_util
@ -136,7 +137,24 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
"""Update the values of the controller."""
update_items(controller, async_add_entities, tracked)
async_dispatcher_connect(hass, controller.event_update, update_controller)
async_dispatcher_connect(hass, controller.signal_update, update_controller)
@callback
def update_disable_on_entities():
"""Update the values of the controller."""
for entity in tracked.values():
disabled_by = None
if not entity.entity_registry_enabled_default and entity.enabled:
disabled_by = DISABLED_CONFIG_ENTRY
registry.async_update_entity(
entity.registry_entry.entity_id, disabled_by=disabled_by
)
async_dispatcher_connect(
hass, controller.signal_options_update, update_disable_on_entities
)
update_controller()
@ -146,65 +164,20 @@ def update_items(controller, async_add_entities, tracked):
"""Update tracked device state from the controller."""
new_tracked = []
if controller.option_track_clients:
for items, tracker_class in (
(controller.api.clients, UniFiClientTracker),
(controller.api.devices, UniFiDeviceTracker),
):
for client_id in controller.api.clients:
for item_id in items:
if client_id in tracked:
if not tracked[client_id].enabled:
continue
LOGGER.debug(
"Updating UniFi tracked client %s (%s)",
tracked[client_id].entity_id,
tracked[client_id].client.mac,
)
tracked[client_id].async_schedule_update_ha_state()
if item_id in tracked:
if tracked[item_id].enabled:
tracked[item_id].async_schedule_update_ha_state()
continue
client = controller.api.clients[client_id]
if (
not client.is_wired
and controller.option_ssid_filter
and client.essid not in controller.option_ssid_filter
):
continue
if not controller.option_track_wired_clients and client.is_wired:
continue
tracked[client_id] = UniFiClientTracker(client, controller)
new_tracked.append(tracked[client_id])
LOGGER.debug(
"New UniFi client tracker %s (%s)",
client.name or client.hostname,
client.mac,
)
if controller.option_track_devices:
for device_id in controller.api.devices:
if device_id in tracked:
if not tracked[device_id].enabled:
continue
LOGGER.debug(
"Updating UniFi tracked device %s (%s)",
tracked[device_id].entity_id,
tracked[device_id].device.mac,
)
tracked[device_id].async_schedule_update_ha_state()
continue
device = controller.api.devices[device_id]
tracked[device_id] = UniFiDeviceTracker(device, controller)
new_tracked.append(tracked[device_id])
LOGGER.debug(
"New UniFi device tracker %s (%s)",
device.name or device.model,
device.mac,
)
tracked[item_id] = tracker_class(items[item_id], controller)
new_tracked.append(tracked[item_id])
if new_tracked:
async_add_entities(new_tracked)
@ -218,8 +191,33 @@ class UniFiClientTracker(ScannerEntity):
self.client = client
self.controller = controller
@property
def entity_registry_enabled_default(self):
"""Return if the entity should be enabled when first added to the entity registry."""
if not self.controller.option_track_clients:
return False
if (
not self.client.is_wired
and self.controller.option_ssid_filter
and self.client.essid not in self.controller.option_ssid_filter
):
return False
if not self.controller.option_track_wired_clients and self.client.is_wired:
return False
return True
async def async_added_to_hass(self):
"""Client entity created."""
LOGGER.debug("New UniFi client tracker %s (%s)", self.name, self.client.mac)
async def async_update(self):
"""Synchronize state with controller."""
LOGGER.debug(
"Updating UniFi tracked client %s (%s)", self.entity_id, self.client.mac
)
await self.controller.request_update()
@property
@ -277,8 +275,23 @@ class UniFiDeviceTracker(ScannerEntity):
self.device = device
self.controller = controller
@property
def entity_registry_enabled_default(self):
"""Return if the entity should be enabled when first added to the entity registry."""
if not self.controller.option_track_devices:
return False
return True
async def async_added_to_hass(self):
"""Subscribe to device events."""
LOGGER.debug("New UniFi device tracker %s (%s)", self.name, self.device.mac)
async def async_update(self):
"""Synchronize state with controller."""
LOGGER.debug(
"Updating UniFi tracked device %s (%s)", self.entity_id, self.device.mac
)
await self.controller.request_update()
@property

4
homeassistant/components/unifi/switch.py
View file

@ -61,7 +61,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
"""Update the values of the controller."""
update_items(controller, async_add_entities, switches, switches_off)
async_dispatcher_connect(hass, controller.event_update, update_controller)
async_dispatcher_connect(hass, controller.signal_update, update_controller)
update_controller()
switches_off.clear()
@ -220,7 +220,7 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
or self.client.sw_mac
and (
self.controller.available
or self.client.sw_mac in self.controller.api.devices
and self.client.sw_mac in self.controller.api.devices
)
)

12
tests/components/unifi/test_controller.py
View file

@ -42,12 +42,12 @@ async def test_controller_setup():
{
CONF_HOST: CONTROLLER_DATA[CONF_HOST],
CONF_SITE_ID: "nice name",
controller.CONF_BLOCK_CLIENT: [],
controller.CONF_TRACK_CLIENTS: True,
controller.CONF_TRACK_DEVICES: True,
controller.CONF_TRACK_WIRED_CLIENTS: True,
controller.CONF_DETECTION_TIME: 300,
controller.CONF_SSID_FILTER: [],
controller.CONF_BLOCK_CLIENT: ["mac"],
controller.CONF_DONT_TRACK_CLIENTS: True,
controller.CONF_DONT_TRACK_DEVICES: True,
controller.CONF_DONT_TRACK_WIRED_CLIENTS: True,
controller.CONF_DETECTION_TIME: 30,
controller.CONF_SSID_FILTER: ["ssid"],
}
]
}

26
tests/components/unifi/test_device_tracker.py
View file

@ -15,6 +15,8 @@ from homeassistant.components.unifi.const import (
CONF_CONTROLLER,
CONF_SITE_ID,
CONF_SSID_FILTER,
CONF_TRACK_DEVICES,
CONF_TRACK_WIRED_CLIENTS,
UNIFI_CONFIG,
)
from homeassistant.const import (
@ -69,10 +71,10 @@ DEVICE_1 = {
"mac": "00:00:00:00:01:01",
"model": "US16P150",
"name": "device_1",
"overheating": False,
"overheating": True,
"state": 1,
"type": "usw",
"upgradable": False,
"upgradable": True,
"version": "4.0.42.10433",
}
DEVICE_2 = {
@ -149,6 +151,7 @@ async def setup_controller(hass, mock_controller, options={}):
system_options={},
options=options,
)
hass.config_entries._entries.append(config_entry)
mock_controller.config_entry = config_entry
await mock_controller.async_update()
@ -230,6 +233,25 @@ async def test_tracked_devices(hass, mock_controller):
device_1 = hass.states.get("device_tracker.device_1")
assert device_1.state == STATE_UNAVAILABLE
mock_controller.config_entry.add_update_listener(
mock_controller.async_options_updated
)
hass.config_entries.async_update_entry(
mock_controller.config_entry,
options={
CONF_SSID_FILTER: [],
CONF_TRACK_WIRED_CLIENTS: False,
CONF_TRACK_DEVICES: False,
},
)
await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1")
assert client_1
client_2 = hass.states.get("device_tracker.wired_client")
assert client_2 is None
device_1 = hass.states.get("device_tracker.device_1")
assert device_1 is None
async def test_restoring_client(hass, mock_controller):
"""Test the update_items function with some clients."""