UniFi - Bandwidth sensors (#27229)

* First draft of UniFi bandwidth sensors

* Clean up

* Add tests for sensors

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"
+                }
             }
         }
     }

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)

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

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
             )

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}"

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"
+                }
             }
         }
     }

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,

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):

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"

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"