diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index fc0d1324f47..b7cd8e8b6a1 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -164,7 +164,7 @@ class UniFiController: WIRELESS_GUEST_CONNECTED, ): self.update_wireless_clients() - else: + elif data.get("clients") or data.get("devices"): async_dispatcher_send(self.hass, self.signal_update) @property @@ -238,13 +238,13 @@ class UniFiController: self.api.start_websocket() - self.config_entry.add_update_listener(self.async_options_updated) + self.config_entry.add_update_listener(self.async_config_entry_updated) return True @staticmethod - async def async_options_updated(hass, entry): - """Triggered by config entry options updates.""" + async def async_config_entry_updated(hass, entry) -> None: + """Handle signals of config entry being updated.""" controller_id = CONTROLLER_ID.format( host=entry.data[CONF_CONTROLLER][CONF_HOST], site=entry.data[CONF_CONTROLLER][CONF_SITE_ID], diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 859a37049b0..5dd5f0c83ae 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -6,10 +6,8 @@ from homeassistant.components.device_tracker.config_entry import ScannerEntity from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER from homeassistant.components.unifi.config_flow import get_controller_from_config_entry from homeassistant.core import callback -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.util.dt as dt_util from .const import ATTR_MANUFACTURER @@ -43,7 +41,11 @@ async def async_setup_entry(hass, config_entry, async_add_entities): controller = get_controller_from_config_entry(hass, config_entry) tracked = {} - registry = await entity_registry.async_get_registry(hass) + option_track_clients = controller.option_track_clients + option_track_devices = controller.option_track_devices + option_track_wired_clients = controller.option_track_wired_clients + + registry = await hass.helpers.entity_registry.async_get_registry() # Restore clients that is not a part of active clients list. for entity in registry.entities.values(): @@ -65,6 +67,12 @@ async def async_setup_entry(hass, config_entry, async_add_entities): @callback def update_controller(): """Update the values of the controller.""" + nonlocal option_track_clients + nonlocal option_track_devices + + if not option_track_clients and not option_track_devices: + return + add_entities(controller, async_add_entities, tracked) controller.listeners.append( @@ -72,24 +80,59 @@ async def async_setup_entry(hass, config_entry, async_add_entities): ) @callback - def update_disable_on_entities(): - """Update the values of the controller.""" - for entity in tracked.values(): + def options_updated(): + """Manage entities affected by config entry options.""" + nonlocal option_track_clients + nonlocal option_track_devices + nonlocal option_track_wired_clients - if entity.entity_registry_enabled_default == entity.enabled: + update = False + remove = set() + + for current_option, config_entry_option, tracker_class in ( + (option_track_clients, controller.option_track_clients, UniFiClientTracker), + (option_track_devices, controller.option_track_devices, UniFiDeviceTracker), + ): + if current_option == config_entry_option: continue - disabled_by = None - if not entity.entity_registry_enabled_default and entity.enabled: - disabled_by = DISABLED_CONFIG_ENTRY + if config_entry_option: + update = True + else: + for mac, entity in tracked.items(): + if isinstance(entity, tracker_class): + remove.add(mac) - registry.async_update_entity( - entity.registry_entry.entity_id, disabled_by=disabled_by - ) + if ( + controller.option_track_clients + and option_track_wired_clients != controller.option_track_wired_clients + ): + + if controller.option_track_wired_clients: + update = True + else: + for mac, entity in tracked.items(): + if isinstance(entity, UniFiClientTracker) and entity.is_wired: + remove.add(mac) + + option_track_clients = controller.option_track_clients + option_track_devices = controller.option_track_devices + option_track_wired_clients = controller.option_track_wired_clients + + for mac in remove: + entity = tracked.pop(mac) + + if registry.async_is_registered(entity.entity_id): + registry.async_remove(entity.entity_id) + + hass.async_create_task(entity.async_remove()) + + if update: + update_controller() controller.listeners.append( async_dispatcher_connect( - hass, controller.signal_options_update, update_disable_on_entities + hass, controller.signal_options_update, options_updated ) ) @@ -101,16 +144,23 @@ def add_entities(controller, async_add_entities, tracked): """Add new tracker entities from the controller.""" new_tracked = [] - for items, tracker_class in ( - (controller.api.clients, UniFiClientTracker), - (controller.api.devices, UniFiDeviceTracker), + for items, tracker_class, track in ( + (controller.api.clients, UniFiClientTracker, controller.option_track_clients), + (controller.api.devices, UniFiDeviceTracker, controller.option_track_devices), ): + if not track: + continue for item_id in items: if item_id in tracked: continue + if tracker_class is UniFiClientTracker and ( + not controller.option_track_wired_clients and items[item_id].is_wired + ): + continue + tracked[item_id] = tracker_class(items[item_id], controller) new_tracked.append(tracked[item_id]) @@ -130,11 +180,12 @@ class UniFiClientTracker(UniFiClient, ScannerEntity): self.wired_bug = dt_util.utcnow() - self.controller.option_detection_time @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 + def is_connected(self): + """Return true if the client is connected to the network. + If connected to unwanted ssid return False. + If is_wired and client.is_wired differ it means that the device is offline and UniFi bug shows device as wired. + """ if ( not self.is_wired and self.controller.option_ssid_filter @@ -142,17 +193,6 @@ class UniFiClientTracker(UniFiClient, ScannerEntity): ): return False - if not self.controller.option_track_wired_clients and self.is_wired: - return False - - return True - - @property - def is_connected(self): - """Return true if the client is connected to the network. - - If is_wired and client.is_wired differ it means that the device is offline and UniFi bug shows device as wired. - """ if self.is_wired != self.client.is_wired: if not self.wired_bug: self.wired_bug = dt_util.utcnow() @@ -202,13 +242,6 @@ class UniFiDeviceTracker(ScannerEntity): self.controller = controller self.listeners = [] - @property - def entity_registry_enabled_default(self): - """Return if the entity should be enabled when first added to the entity registry.""" - if self.controller.option_track_devices: - return True - return False - async def async_added_to_hass(self): """Subscribe to device events.""" LOGGER.debug("New UniFi device tracker %s (%s)", self.name, self.device.mac) diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 860ddf81d7d..942b0ef6779 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -3,9 +3,7 @@ import logging from homeassistant.components.unifi.config_flow import get_controller_from_config_entry from homeassistant.core import callback -from homeassistant.helpers import entity_registry from homeassistant.helpers.dispatcher import async_dispatcher_connect -from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY from .unifi_client import UniFiClient @@ -24,11 +22,18 @@ async def async_setup_entry(hass, config_entry, async_add_entities): controller = get_controller_from_config_entry(hass, config_entry) sensors = {} - registry = await entity_registry.async_get_registry(hass) + option_allow_bandwidth_sensors = controller.option_allow_bandwidth_sensors + + entity_registry = await hass.helpers.entity_registry.async_get_registry() @callback def update_controller(): """Update the values of the controller.""" + nonlocal option_allow_bandwidth_sensors + + if not option_allow_bandwidth_sensors: + return + add_entities(controller, async_add_entities, sensors) controller.listeners.append( @@ -36,24 +41,29 @@ async def async_setup_entry(hass, config_entry, async_add_entities): ) @callback - def update_disable_on_entities(): + def options_updated(): """Update the values of the controller.""" - for entity in sensors.values(): + nonlocal option_allow_bandwidth_sensors - if entity.entity_registry_enabled_default == entity.enabled: - continue + if option_allow_bandwidth_sensors != controller.option_allow_bandwidth_sensors: + option_allow_bandwidth_sensors = controller.option_allow_bandwidth_sensors - disabled_by = None - if not entity.entity_registry_enabled_default and entity.enabled: - disabled_by = DISABLED_CONFIG_ENTRY + if option_allow_bandwidth_sensors: + update_controller() - registry.async_update_entity( - entity.registry_entry.entity_id, disabled_by=disabled_by - ) + else: + for sensor in sensors.values(): + + if entity_registry.async_is_registered(sensor.entity_id): + entity_registry.async_remove(sensor.entity_id) + + hass.async_create_task(sensor.async_remove()) + + sensors.clear() controller.listeners.append( async_dispatcher_connect( - hass, controller.signal_options_update, update_disable_on_entities + hass, controller.signal_options_update, options_updated ) ) @@ -87,13 +97,6 @@ def add_entities(controller, async_add_entities, sensors): class UniFiRxBandwidthSensor(UniFiClient): """Receiving bandwidth sensor.""" - @property - def entity_registry_enabled_default(self): - """Return if the entity should be enabled when first added to the entity registry.""" - if self.controller.option_allow_bandwidth_sensors: - return True - return False - @property def state(self): """Return the state of the sensor.""" diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 86595fe43e2..608e72b483a 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -11,6 +11,7 @@ from homeassistant.components import unifi import homeassistant.components.device_tracker as device_tracker from homeassistant.components.unifi.const import ( CONF_SSID_FILTER, + CONF_TRACK_CLIENTS, CONF_TRACK_DEVICES, CONF_TRACK_WIRED_CLIENTS, ) @@ -114,7 +115,7 @@ async def test_tracked_devices(hass): devices_response=[DEVICE_1, DEVICE_2], known_wireless_clients=(CLIENT_4["mac"],), ) - assert len(hass.states.async_all()) == 5 + assert len(hass.states.async_all()) == 6 client_1 = hass.states.get("device_tracker.client_1") assert client_1 is not None @@ -125,7 +126,8 @@ async def test_tracked_devices(hass): assert client_2.state == "not_home" client_3 = hass.states.get("device_tracker.client_3") - assert client_3 is None + assert client_3 is not None + assert client_3.state == "not_home" # Wireless client with wired bug, if bug active on restart mark device away client_4 = hass.states.get("device_tracker.client_4") @@ -136,6 +138,7 @@ async def test_tracked_devices(hass): assert device_1 is not None assert device_1.state == "not_home" + # State change signalling works client_1_copy = copy(CLIENT_1) client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) event = {"meta": {"message": "sta:sync"}, "data": [client_1_copy]} @@ -152,28 +155,6 @@ async def test_tracked_devices(hass): device_1 = hass.states.get("device_tracker.device_1") assert device_1.state == "home" - # Controller unavailable - controller.async_unifi_signalling_callback( - SIGNAL_CONNECTION_STATE, STATE_DISCONNECTED - ) - await hass.async_block_till_done() - - client_1 = hass.states.get("device_tracker.client_1") - assert client_1.state == STATE_UNAVAILABLE - - device_1 = hass.states.get("device_tracker.device_1") - assert device_1.state == STATE_UNAVAILABLE - - # Controller available - controller.async_unifi_signalling_callback(SIGNAL_CONNECTION_STATE, STATE_RUNNING) - await hass.async_block_till_done() - - client_1 = hass.states.get("device_tracker.client_1") - assert client_1.state == "home" - - device_1 = hass.states.get("device_tracker.device_1") - assert device_1.state == "home" - # Disabled device is unavailable device_1_copy = copy(DEVICE_1) device_1_copy["disabled"] = True @@ -184,24 +165,205 @@ async def test_tracked_devices(hass): device_1 = hass.states.get("device_tracker.device_1") assert device_1.state == STATE_UNAVAILABLE - # Don't track wired clients nor devices - controller.config_entry.add_update_listener(controller.async_options_updated) - hass.config_entries.async_update_entry( - controller.config_entry, - options={ - CONF_SSID_FILTER: [], - CONF_TRACK_WIRED_CLIENTS: False, - CONF_TRACK_DEVICES: False, - }, + +async def test_controller_state_change(hass): + """Verify entities state reflect on controller becoming unavailable.""" + controller = await setup_unifi_integration( + hass, clients_response=[CLIENT_1], devices_response=[DEVICE_1], + ) + assert len(hass.states.async_all()) == 3 + + # Controller unavailable + controller.async_unifi_signalling_callback( + SIGNAL_CONNECTION_STATE, STATE_DISCONNECTED ) await hass.async_block_till_done() + client_1 = hass.states.get("device_tracker.client_1") - assert client_1 + assert client_1.state == STATE_UNAVAILABLE + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1.state == STATE_UNAVAILABLE + + # Controller available + controller.async_unifi_signalling_callback(SIGNAL_CONNECTION_STATE, STATE_RUNNING) + await hass.async_block_till_done() + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1.state == "not_home" + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1.state == "not_home" + + +async def test_option_track_clients(hass): + """Test the tracking of clients can be turned off.""" + controller = await setup_unifi_integration( + hass, clients_response=[CLIENT_1, CLIENT_2], devices_response=[DEVICE_1], + ) + assert len(hass.states.async_all()) == 4 + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is not None + + client_2 = hass.states.get("device_tracker.wired_client") + assert client_2 is not None + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1 is not None + + hass.config_entries.async_update_entry( + controller.config_entry, options={CONF_TRACK_CLIENTS: False}, + ) + await hass.async_block_till_done() + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is None + 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 not None + + hass.config_entries.async_update_entry( + controller.config_entry, options={CONF_TRACK_CLIENTS: True}, + ) + await hass.async_block_till_done() + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is not None + + client_2 = hass.states.get("device_tracker.wired_client") + assert client_2 is not None + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1 is not None + + +async def test_option_track_wired_clients(hass): + """Test the tracking of wired clients can be turned off.""" + controller = await setup_unifi_integration( + hass, clients_response=[CLIENT_1, CLIENT_2], devices_response=[DEVICE_1], + ) + assert len(hass.states.async_all()) == 4 + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is not None + + client_2 = hass.states.get("device_tracker.wired_client") + assert client_2 is not None + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1 is not None + + hass.config_entries.async_update_entry( + controller.config_entry, options={CONF_TRACK_WIRED_CLIENTS: False}, + ) + await hass.async_block_till_done() + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is not None + + 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 not None + + hass.config_entries.async_update_entry( + controller.config_entry, options={CONF_TRACK_WIRED_CLIENTS: True}, + ) + await hass.async_block_till_done() + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is not None + + client_2 = hass.states.get("device_tracker.wired_client") + assert client_2 is not None + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1 is not None + + +async def test_option_track_devices(hass): + """Test the tracking of devices can be turned off.""" + controller = await setup_unifi_integration( + hass, clients_response=[CLIENT_1, CLIENT_2], devices_response=[DEVICE_1], + ) + assert len(hass.states.async_all()) == 4 + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is not None + + client_2 = hass.states.get("device_tracker.wired_client") + assert client_2 is not None + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1 is not None + + hass.config_entries.async_update_entry( + controller.config_entry, options={CONF_TRACK_DEVICES: False}, + ) + await hass.async_block_till_done() + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is not None + + client_2 = hass.states.get("device_tracker.wired_client") + assert client_2 is not None + device_1 = hass.states.get("device_tracker.device_1") assert device_1 is None + hass.config_entries.async_update_entry( + controller.config_entry, options={CONF_TRACK_DEVICES: True}, + ) + await hass.async_block_till_done() + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is not None + + client_2 = hass.states.get("device_tracker.wired_client") + assert client_2 is not None + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1 is not None + + +async def test_option_ssid_filter(hass): + """Test the SSID filter works.""" + controller = await setup_unifi_integration( + hass, options={CONF_SSID_FILTER: ["ssid"]}, clients_response=[CLIENT_3], + ) + assert len(hass.states.async_all()) == 2 + + # SSID filter active + client_3 = hass.states.get("device_tracker.client_3") + assert client_3.state == "not_home" + + client_3_copy = copy(CLIENT_3) + client_3_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) + event = {"meta": {"message": "sta:sync"}, "data": [client_3_copy]} + controller.api.message_handler(event) + await hass.async_block_till_done() + + # SSID filter active even though time stamp should mark as home + client_3 = hass.states.get("device_tracker.client_3") + assert client_3.state == "not_home" + + # Remove SSID filter + hass.config_entries.async_update_entry( + controller.config_entry, options={CONF_SSID_FILTER: []}, + ) + event = {"meta": {"message": "sta:sync"}, "data": [client_3_copy]} + controller.api.message_handler(event) + await hass.async_block_till_done() + + # SSID no longer filtered + client_3 = hass.states.get("device_tracker.client_3") + assert client_3.state == "home" + async def test_wireless_client_go_wired_issue(hass): """Test the solution to catch wireless device go wired UniFi issue. diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index 668b7a36ada..c726d3bb1cb 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -99,3 +99,27 @@ async def test_sensors(hass): wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx") assert wireless_client_tx.state == "6789.0" + + hass.config_entries.async_update_entry( + controller.config_entry, + options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: False}, + ) + await hass.async_block_till_done() + + wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx") + assert wireless_client_rx is None + + wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx") + assert wireless_client_tx is None + + hass.config_entries.async_update_entry( + controller.config_entry, + options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True}, + ) + await hass.async_block_till_done() + + wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx") + assert wireless_client_rx.state == "2345.0" + + wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx") + assert wireless_client_tx.state == "6789.0"