Add AsusWRT Devices Connected Sensor (#34204)

* Add Devices Connected

* Remove attributes

* Add sensors to test

* Improve sensor test

* Cleanup

* Apply suggestions from code review

Co-Authored-By: Martin Hjelmare <marhje52@gmail.com>

* Import device from aioasuswrt

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/asuswrt/__init__.py

@@ -36,7 +36,7 @@ FIRST_RETRY_TIME = 60
 MAX_RETRY_TIME = 900
 
 SECRET_GROUP = "Password or SSH Key"
-SENSOR_TYPES = ["upload_speed", "download_speed", "download", "upload"]
+SENSOR_TYPES = ["devices", "upload_speed", "download_speed", "download", "upload"]
 
 CONFIG_SCHEMA = vol.Schema(
     {

homeassistant/components/asuswrt/sensor.py

@@ -1,6 +1,8 @@
 """Asuswrt status sensors."""
 import logging
 
+from aioasuswrt.asuswrt import AsusWrt
+
 from homeassistant.const import DATA_GIGABYTES, DATA_RATE_MEGABITS_PER_SECOND
 from homeassistant.helpers.entity import Entity
 
@@ -18,6 +20,8 @@ async def async_setup_platform(hass, config, add_entities, discovery_info=None):
 
     devices = []
 
+    if "devices" in discovery_info:
+        devices.append(AsuswrtDevicesSensor(api))
     if "download" in discovery_info:
         devices.append(AsuswrtTotalRXSensor(api))
     if "upload" in discovery_info:
@@ -35,10 +39,11 @@ class AsuswrtSensor(Entity):
 
     _name = "generic"
 
-    def __init__(self, api):
+    def __init__(self, api: AsusWrt):
         """Initialize the sensor."""
         self._api = api
         self._state = None
+        self._devices = None
         self._rates = None
         self._speed = None
 
@@ -54,10 +59,23 @@ class AsuswrtSensor(Entity):
 
     async def async_update(self):
         """Fetch status from asuswrt."""
+        self._devices = await self._api.async_get_connected_devices()
         self._rates = await self._api.async_get_bytes_total()
         self._speed = await self._api.async_get_current_transfer_rates()
 
 
+class AsuswrtDevicesSensor(AsuswrtSensor):
+    """Representation of a asuswrt download speed sensor."""
+
+    _name = "Asuswrt Devices Connected"
+
+    async def async_update(self):
+        """Fetch new state data for the sensor."""
+        await super().async_update()
+        if self._devices:
+            self._state = len(self._devices)
+
+
 class AsuswrtRXSensor(AsuswrtSensor):
     """Representation of a asuswrt download speed sensor."""
 

tests/components/asuswrt/test_sensor.py

@@ -1,7 +1,10 @@
-"""The tests for the ASUSWRT sensor platform."""
-from unittest.mock import patch
+"""The tests for the AsusWrt sensor platform."""
+from datetime import datetime, timedelta
 
-# import homeassistant.components.sensor as sensor
+from aioasuswrt.asuswrt import Device
+from asynctest import CoroutineMock, patch
+
+from homeassistant.components import sensor
 from homeassistant.components.asuswrt import (
     CONF_DNSMASQ,
     CONF_INTERFACE,
@@ -9,13 +12,15 @@ from homeassistant.components.asuswrt import (
     CONF_PORT,
     CONF_PROTOCOL,
     CONF_SENSORS,
-    DATA_ASUSWRT,
     DOMAIN,
 )
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
+from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
+import homeassistant.util.dt as dt_util
+from homeassistant.util.dt import utcnow
 
-from tests.common import mock_coro_func
+from tests.common import async_fire_time_changed
 
 VALID_CONFIG_ROUTER_SSH = {
     DOMAIN: {
@@ -27,16 +32,52 @@ VALID_CONFIG_ROUTER_SSH = {
         CONF_PROTOCOL: "ssh",
         CONF_USERNAME: "fake_user",
         CONF_PASSWORD: "fake_pass",
-        CONF_SENSORS: "upload",
+        CONF_SENSORS: [
+            "devices",
+            "download_speed",
+            "download",
+            "upload_speed",
+            "upload",
+        ],
     }
 }
 
+MOCK_DEVICES = {
+    "a1:b1:c1:d1:e1:f1": Device("a1:b1:c1:d1:e1:f1", "192.168.1.2", "Test"),
+    "a2:b2:c2:d2:e2:f2": Device("a2:b2:c2:d2:e2:f2", "192.168.1.3", "TestTwo"),
+    "a3:b3:c3:d3:e3:f3": Device("a3:b3:c3:d3:e3:f3", "192.168.1.4", "TestThree"),
+}
+MOCK_BYTES_TOTAL = [60000000000, 50000000000]
+MOCK_CURRENT_TRANSFER_RATES = [20000000, 10000000]
 
-async def test_default_sensor_setup(hass):
+
+async def test_sensors(hass: HomeAssistant):
     """Test creating an AsusWRT sensor."""
     with patch("homeassistant.components.asuswrt.AsusWrt") as AsusWrt:
-        AsusWrt().connection.async_connect = mock_coro_func()
+        AsusWrt().connection.async_connect = CoroutineMock()
+        AsusWrt().async_get_connected_devices = CoroutineMock(return_value=MOCK_DEVICES)
+        AsusWrt().async_get_bytes_total = CoroutineMock(return_value=MOCK_BYTES_TOTAL)
+        AsusWrt().async_get_current_transfer_rates = CoroutineMock(
+            return_value=MOCK_CURRENT_TRANSFER_RATES
+        )
 
-        result = await async_setup_component(hass, DOMAIN, VALID_CONFIG_ROUTER_SSH)
-        assert result
-        assert hass.data[DATA_ASUSWRT] is not None
+        now = datetime(2020, 1, 1, 1, tzinfo=dt_util.UTC)
+        with patch(("homeassistant.helpers.event.dt_util.utcnow"), return_value=now):
+            assert await async_setup_component(hass, DOMAIN, VALID_CONFIG_ROUTER_SSH)
+            await hass.async_block_till_done()
+            async_fire_time_changed(hass, utcnow() + timedelta(seconds=30))
+            await hass.async_block_till_done()
+
+            assert (
+                hass.states.get(f"{sensor.DOMAIN}.asuswrt_devices_connected").state
+                == "3"
+            )
+            assert (
+                hass.states.get(f"{sensor.DOMAIN}.asuswrt_download_speed").state
+                == "160.0"
+            )
+            assert hass.states.get(f"{sensor.DOMAIN}.asuswrt_download").state == "60.0"
+            assert (
+                hass.states.get(f"{sensor.DOMAIN}.asuswrt_upload_speed").state == "80.0"
+            )
+            assert hass.states.get(f"{sensor.DOMAIN}.asuswrt_upload").state == "50.0"