diff --git a/homeassistant/components/asuswrt/__init__.py b/homeassistant/components/asuswrt/__init__.py index 446fe898aaa..d9c87beea5d 100644 --- a/homeassistant/components/asuswrt/__init__.py +++ b/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( { diff --git a/homeassistant/components/asuswrt/sensor.py b/homeassistant/components/asuswrt/sensor.py index 50100d3625d..631e6e9d70f 100644 --- a/homeassistant/components/asuswrt/sensor.py +++ b/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.""" diff --git a/tests/components/asuswrt/test_sensor.py b/tests/components/asuswrt/test_sensor.py index 39443c3fef8..991be0bac50 100644 --- a/tests/components/asuswrt/test_sensor.py +++ b/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"