diff --git a/homeassistant/components/unifi/.translations/en.json b/homeassistant/components/unifi/.translations/en.json index 2025bad6246..d9b65b6d1da 100644 --- a/homeassistant/components/unifi/.translations/en.json +++ b/homeassistant/components/unifi/.translations/en.json @@ -32,6 +32,11 @@ "track_devices": "Track network devices (Ubiquiti devices)", "track_wired_clients": "Include wired network clients" } + }, + "statistics_sensors": { + "data": { + "allow_bandwidth_sensors": "Create bandwidth usage sensors for network clients" + } } } } diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index fdb75d09194..92281837f48 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -12,6 +12,7 @@ from homeassistant.const import ( ) from .const import ( + CONF_ALLOW_BANDWIDTH_SENSORS, CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_SITE_ID, @@ -19,6 +20,7 @@ from .const import ( CONF_TRACK_DEVICES, CONF_TRACK_WIRED_CLIENTS, CONTROLLER_ID, + DEFAULT_ALLOW_BANDWIDTH_SENSORS, DEFAULT_TRACK_CLIENTS, DEFAULT_TRACK_DEVICES, DEFAULT_TRACK_WIRED_CLIENTS, @@ -171,6 +173,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow): def __init__(self, config_entry): """Initialize UniFi options flow.""" self.config_entry = config_entry + self.options = dict(config_entry.options) async def async_step_init(self, user_input=None): """Manage the UniFi options.""" @@ -179,7 +182,8 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow): async def async_step_device_tracker(self, user_input=None): """Manage the device tracker options.""" if user_input is not None: - return self.async_create_entry(title="", data=user_input) + self.options.update(user_input) + return await self.async_step_statistics_sensors() return self.async_show_form( step_id="device_tracker", @@ -212,3 +216,28 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow): } ), ) + + async def async_step_statistics_sensors(self, user_input=None): + """Manage the statistics sensors options.""" + if user_input is not None: + self.options.update(user_input) + return await self._update_options() + + return self.async_show_form( + step_id="statistics_sensors", + data_schema=vol.Schema( + { + vol.Optional( + CONF_ALLOW_BANDWIDTH_SENSORS, + default=self.config_entry.options.get( + CONF_ALLOW_BANDWIDTH_SENSORS, + DEFAULT_ALLOW_BANDWIDTH_SENSORS, + ), + ): bool + } + ), + ) + + async def _update_options(self): + """Update config entry options.""" + return self.async_create_entry(title="", data=self.options) diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py index eac14735074..d82b7b49d45 100644 --- a/homeassistant/components/unifi/const.py +++ b/homeassistant/components/unifi/const.py @@ -12,6 +12,7 @@ CONF_SITE_ID = "site" UNIFI_CONFIG = "unifi_config" UNIFI_WIRELESS_CLIENTS = "unifi_wireless_clients" +CONF_ALLOW_BANDWIDTH_SENSORS = "allow_bandwidth_sensors" CONF_BLOCK_CLIENT = "block_client" CONF_DETECTION_TIME = "detection_time" CONF_TRACK_CLIENTS = "track_clients" @@ -23,6 +24,7 @@ CONF_DONT_TRACK_CLIENTS = "dont_track_clients" CONF_DONT_TRACK_DEVICES = "dont_track_devices" CONF_DONT_TRACK_WIRED_CLIENTS = "dont_track_wired_clients" +DEFAULT_ALLOW_BANDWIDTH_SENSORS = False DEFAULT_BLOCK_CLIENTS = [] DEFAULT_TRACK_CLIENTS = True DEFAULT_TRACK_DEVICES = True diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index ffea98b9050..fa1164166bd 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -15,6 +15,7 @@ from homeassistant.helpers import aiohttp_client from homeassistant.helpers.dispatcher import async_dispatcher_send from .const import ( + CONF_ALLOW_BANDWIDTH_SENSORS, CONF_BLOCK_CLIENT, CONF_CONTROLLER, CONF_DETECTION_TIME, @@ -27,6 +28,7 @@ from .const import ( CONF_SITE_ID, CONF_SSID_FILTER, CONTROLLER_ID, + DEFAULT_ALLOW_BANDWIDTH_SENSORS, DEFAULT_BLOCK_CLIENTS, DEFAULT_TRACK_CLIENTS, DEFAULT_TRACK_DEVICES, @@ -40,6 +42,8 @@ from .const import ( ) from .errors import AuthenticationRequired, CannotConnect +SUPPORTED_PLATFORMS = ["device_tracker", "sensor", "switch"] + class UniFiController: """Manages a single UniFi Controller.""" @@ -76,6 +80,13 @@ class UniFiController: """Return the site user role of this controller.""" return self._site_role + @property + def option_allow_bandwidth_sensors(self): + """Config entry option to allow bandwidth sensors.""" + return self.config_entry.options.get( + CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS + ) + @property def option_block_clients(self): """Config entry option with list of clients to control network access.""" @@ -225,7 +236,7 @@ class UniFiController: self.config_entry.add_update_listener(self.async_options_updated) - for platform in ["device_tracker", "switch"]: + for platform in SUPPORTED_PLATFORMS: hass.async_create_task( hass.config_entries.async_forward_entry_setup( self.config_entry, platform @@ -294,7 +305,7 @@ class UniFiController: if self.api is None: return True - for platform in ["device_tracker", "switch"]: + for platform in SUPPORTED_PLATFORMS: await self.hass.config_entries.async_forward_entry_unload( self.config_entry, platform ) diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py new file mode 100644 index 00000000000..aad013970d1 --- /dev/null +++ b/homeassistant/components/unifi/sensor.py @@ -0,0 +1,168 @@ +"""Support for bandwidth sensors with UniFi clients.""" +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.device_registry import CONNECTION_NETWORK_MAC +from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity import Entity +from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY + +LOGGER = logging.getLogger(__name__) + +ATTR_RECEIVING = "receiving" +ATTR_TRANSMITTING = "transmitting" + + +async def async_setup_platform(hass, config, async_add_entities, discovery_info=None): + """Sensor platform doesn't support configuration through configuration.yaml.""" + + +async def async_setup_entry(hass, config_entry, async_add_entities): + """Set up sensors for UniFi integration.""" + controller = get_controller_from_config_entry(hass, config_entry) + sensors = {} + + registry = await entity_registry.async_get_registry(hass) + + @callback + def update_controller(): + """Update the values of the controller.""" + update_items(controller, async_add_entities, sensors) + + async_dispatcher_connect(hass, controller.signal_update, update_controller) + + @callback + def update_disable_on_entities(): + """Update the values of the controller.""" + for entity in sensors.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() + + +@callback +def update_items(controller, async_add_entities, sensors): + """Update sensors from the controller.""" + new_sensors = [] + + for client_id in controller.api.clients: + for direction, sensor_class in ( + ("rx", UniFiRxBandwidthSensor), + ("tx", UniFiTxBandwidthSensor), + ): + item_id = f"{direction}-{client_id}" + + if item_id in sensors: + sensor = sensors[item_id] + if sensor.enabled: + sensor.async_schedule_update_ha_state() + continue + + sensors[item_id] = sensor_class( + controller.api.clients[client_id], controller + ) + new_sensors.append(sensors[item_id]) + + if new_sensors: + async_add_entities(new_sensors) + + +class UniFiBandwidthSensor(Entity): + """UniFi Bandwidth sensor base class.""" + + def __init__(self, client, controller): + """Set up client.""" + self.client = client + self.controller = controller + self.is_wired = self.client.mac not in controller.wireless_clients + + @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 + + async def async_added_to_hass(self): + """Client entity created.""" + LOGGER.debug("New UniFi bandwidth sensor %s (%s)", self.name, self.client.mac) + + async def async_update(self): + """Synchronize state with controller. + + Make sure to update self.is_wired if client is wireless, there is an issue when clients go offline that they get marked as wired. + """ + LOGGER.debug( + "Updating UniFi bandwidth sensor %s (%s)", self.entity_id, self.client.mac + ) + await self.controller.request_update() + + if self.is_wired and self.client.mac in self.controller.wireless_clients: + self.is_wired = False + + @property + def available(self) -> bool: + """Return if controller is available.""" + return self.controller.available + + @property + def device_info(self): + """Return a device description for device registry.""" + return {"connections": {(CONNECTION_NETWORK_MAC, self.client.mac)}} + + +class UniFiRxBandwidthSensor(UniFiBandwidthSensor): + """Receiving bandwidth sensor.""" + + @property + def state(self): + """Return the state of the sensor.""" + if self.is_wired: + return self.client.wired_rx_bytes / 1000000 + return self.client.raw.get("rx_bytes", 0) / 1000000 + + @property + def name(self): + """Return the name of the client.""" + name = self.client.name or self.client.hostname + return f"{name} RX" + + @property + def unique_id(self): + """Return a unique identifier for this bandwidth sensor.""" + return f"rx-{self.client.mac}" + + +class UniFiTxBandwidthSensor(UniFiBandwidthSensor): + """Transmitting bandwidth sensor.""" + + @property + def state(self): + """Return the state of the sensor.""" + if self.is_wired: + return self.client.wired_tx_bytes / 1000000 + return self.client.raw.get("tx_bytes", 0) / 1000000 + + @property + def name(self): + """Return the name of the client.""" + name = self.client.name or self.client.hostname + return f"{name} TX" + + @property + def unique_id(self): + """Return a unique identifier for this bandwidth sensor.""" + return f"tx-{self.client.mac}" diff --git a/homeassistant/components/unifi/strings.json b/homeassistant/components/unifi/strings.json index c484bfbf09f..ce2f2345917 100644 --- a/homeassistant/components/unifi/strings.json +++ b/homeassistant/components/unifi/strings.json @@ -35,6 +35,11 @@ "track_devices": "Track network devices (Ubiquiti devices)", "track_wired_clients": "Include wired network clients" } + }, + "statistics_sensors": { + "data": { + "allow_bandwidth_sensors": "Create bandwidth usage sensors for network clients" + } } } } diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index f0183a7ecb3..f8fad6dac8e 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -14,7 +14,6 @@ LOGGER = logging.getLogger(__name__) async def async_setup_platform(hass, config, async_add_entities, discovery_info=None): """Component doesn't support configuration through configuration.yaml.""" - pass async def async_setup_entry(hass, config_entry, async_add_entities): @@ -231,8 +230,6 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity): """Return the device state attributes.""" attributes = { "power": self.port.poe_power, - "received": self.client.wired_rx_bytes / 1000000, - "sent": self.client.wired_tx_bytes / 1000000, "switch": self.client.sw_mac, "port": self.client.sw_port, "poe_mode": self.poe_mode, diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index e73719205f7..ae6f3776b4f 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -67,12 +67,18 @@ async def test_controller_setup(): assert await unifi_controller.async_setup() is True assert unifi_controller.api is api - assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == 2 + assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == len( + controller.SUPPORTED_PLATFORMS + ) assert hass.config_entries.async_forward_entry_setup.mock_calls[0][1] == ( entry, "device_tracker", ) assert hass.config_entries.async_forward_entry_setup.mock_calls[1][1] == ( + entry, + "sensor", + ) + assert hass.config_entries.async_forward_entry_setup.mock_calls[2][1] == ( entry, "switch", ) @@ -214,12 +220,16 @@ async def test_reset_unloads_entry_if_setup(): with patch.object(controller, "get_controller", return_value=mock_coro(api)): assert await unifi_controller.async_setup() is True - assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == 2 + assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == len( + controller.SUPPORTED_PLATFORMS + ) hass.config_entries.async_forward_entry_unload.return_value = mock_coro(True) assert await unifi_controller.async_reset() - assert len(hass.config_entries.async_forward_entry_unload.mock_calls) == 2 + assert len(hass.config_entries.async_forward_entry_unload.mock_calls) == len( + controller.SUPPORTED_PLATFORMS + ) async def test_get_controller(hass): diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py new file mode 100644 index 00000000000..9064f1c9aba --- /dev/null +++ b/tests/components/unifi/test_sensor.py @@ -0,0 +1,207 @@ +"""UniFi sensor platform tests.""" +from collections import deque +from copy import deepcopy + +from asynctest import patch + +from homeassistant import config_entries +from homeassistant.components import unifi +from homeassistant.components.unifi.const import ( + CONF_CONTROLLER, + CONF_SITE_ID, + CONTROLLER_ID as CONF_CONTROLLER_ID, + UNIFI_CONFIG, + UNIFI_WIRELESS_CLIENTS, +) +from homeassistant.setup import async_setup_component +from homeassistant.const import ( + CONF_HOST, + CONF_PASSWORD, + CONF_PORT, + CONF_USERNAME, + CONF_VERIFY_SSL, +) + +import homeassistant.components.sensor as sensor + +CLIENTS = [ + { + "hostname": "Wired client hostname", + "ip": "10.0.0.1", + "is_wired": True, + "last_seen": 1562600145, + "mac": "00:00:00:00:00:01", + "name": "Wired client name", + "oui": "Producer", + "sw_mac": "00:00:00:00:01:01", + "sw_port": 1, + "wired-rx_bytes": 1234000000, + "wired-tx_bytes": 5678000000, + }, + { + "hostname": "Wireless client hostname", + "ip": "10.0.0.2", + "is_wired": False, + "last_seen": 1562600145, + "mac": "00:00:00:00:00:02", + "name": "Wireless client name", + "oui": "Producer", + "sw_mac": "00:00:00:00:01:01", + "sw_port": 2, + "rx_bytes": 1234000000, + "tx_bytes": 5678000000, + }, +] + +CONTROLLER_DATA = { + CONF_HOST: "mock-host", + CONF_USERNAME: "mock-user", + CONF_PASSWORD: "mock-pswd", + CONF_PORT: 1234, + CONF_SITE_ID: "mock-site", + CONF_VERIFY_SSL: False, +} + +ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA} + +CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site") + +SITES = {"Site name": {"desc": "Site name", "name": "mock-site", "role": "admin"}} + + +async def setup_unifi_integration( + hass, + config, + options, + sites, + clients_response, + devices_response, + clients_all_response, +): + """Create the UniFi controller.""" + hass.data[UNIFI_CONFIG] = [] + hass.data[UNIFI_WIRELESS_CLIENTS] = unifi.UnifiWirelessClients(hass) + config_entry = config_entries.ConfigEntry( + version=1, + domain=unifi.DOMAIN, + title="Mock Title", + data=config, + source="test", + connection_class=config_entries.CONN_CLASS_LOCAL_POLL, + system_options={}, + options=options, + entry_id=1, + ) + + mock_client_responses = deque() + mock_client_responses.append(clients_response) + + mock_device_responses = deque() + mock_device_responses.append(devices_response) + + mock_client_all_responses = deque() + mock_client_all_responses.append(clients_all_response) + + mock_requests = [] + + async def mock_request(self, method, path, json=None): + mock_requests.append({"method": method, "path": path, "json": json}) + + if path == "s/{site}/stat/sta" and mock_client_responses: + return mock_client_responses.popleft() + if path == "s/{site}/stat/device" and mock_device_responses: + return mock_device_responses.popleft() + if path == "s/{site}/rest/user" and mock_client_all_responses: + return mock_client_all_responses.popleft() + return {} + + with patch("aiounifi.Controller.login", return_value=True), patch( + "aiounifi.Controller.sites", return_value=sites + ), patch("aiounifi.Controller.request", new=mock_request): + await unifi.async_setup_entry(hass, config_entry) + await hass.async_block_till_done() + hass.config_entries._entries.append(config_entry) + + controller_id = unifi.get_controller_id_from_config_entry(config_entry) + controller = hass.data[unifi.DOMAIN][controller_id] + + controller.mock_client_responses = mock_client_responses + controller.mock_device_responses = mock_device_responses + controller.mock_client_all_responses = mock_client_all_responses + controller.mock_requests = mock_requests + + return controller + + +async def test_platform_manually_configured(hass): + """Test that we do not discover anything or try to set up a controller.""" + assert ( + await async_setup_component( + hass, sensor.DOMAIN, {sensor.DOMAIN: {"platform": "unifi"}} + ) + is True + ) + assert unifi.DOMAIN not in hass.data + + +async def test_no_clients(hass): + """Test the update_clients function when no clients are found.""" + controller = await setup_unifi_integration( + hass, + ENTRY_CONFIG, + options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True}, + sites=SITES, + clients_response=[], + devices_response=[], + clients_all_response=[], + ) + + assert len(controller.mock_requests) == 3 + assert len(hass.states.async_all()) == 2 + + +async def test_switches(hass): + """Test the update_items function with some clients.""" + controller = await setup_unifi_integration( + hass, + ENTRY_CONFIG, + options={ + unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True, + unifi.const.CONF_TRACK_CLIENTS: False, + unifi.const.CONF_TRACK_DEVICES: False, + }, + sites=SITES, + clients_response=CLIENTS, + devices_response=[], + clients_all_response=[], + ) + + assert len(controller.mock_requests) == 3 + assert len(hass.states.async_all()) == 6 + + wired_client_rx = hass.states.get("sensor.wired_client_name_rx") + assert wired_client_rx.state == "1234.0" + + wired_client_tx = hass.states.get("sensor.wired_client_name_tx") + assert wired_client_tx.state == "5678.0" + + wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx") + assert wireless_client_rx.state == "1234.0" + + wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx") + assert wireless_client_tx.state == "5678.0" + + clients = deepcopy(CLIENTS) + clients[0]["is_wired"] = False + clients[1]["rx_bytes"] = 2345000000 + clients[1]["tx_bytes"] = 6789000000 + + controller.mock_client_responses.append(clients) + await controller.async_update() + 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" diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 56a96b2b5b2..97dda441527 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -264,7 +264,7 @@ async def setup_unifi_integration( async def mock_request(self, method, path, json=None): mock_requests.append({"method": method, "path": path, "json": json}) - print(mock_requests, mock_client_responses, mock_device_responses) + if path == "s/{site}/stat/sta" and mock_client_responses: return mock_client_responses.popleft() if path == "s/{site}/stat/device" and mock_device_responses: @@ -386,8 +386,6 @@ async def test_switches(hass): assert switch_1 is not None assert switch_1.state == "on" assert switch_1.attributes["power"] == "2.56" - assert switch_1.attributes["received"] == 1234 - assert switch_1.attributes["sent"] == 5678 assert switch_1.attributes["switch"] == "00:00:00:00:01:01" assert switch_1.attributes["port"] == 1 assert switch_1.attributes["poe_mode"] == "auto"