From 47f6313e5bd53dda97f1d1c99daa8f6dfbc93833 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 26 Oct 2021 20:23:20 +0200 Subject: [PATCH] Add entity category to UniFi sensors and switches (#58484) --- homeassistant/components/unifi/sensor.py | 4 +++- homeassistant/components/unifi/switch.py | 7 +++++++ tests/components/unifi/test_sensor.py | 14 ++++++++++++++ tests/components/unifi/test_switch.py | 9 +++++++++ 4 files changed, 33 insertions(+), 1 deletion(-) diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 5cbf61d9635..a9e89876459 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -7,7 +7,7 @@ Support for uptime sensors of network clients. from datetime import datetime, timedelta from homeassistant.components.sensor import DEVICE_CLASS_TIMESTAMP, DOMAIN, SensorEntity -from homeassistant.const import DATA_MEGABYTES +from homeassistant.const import DATA_MEGABYTES, ENTITY_CATEGORY_DIAGNOSTIC from homeassistant.core import callback from homeassistant.helpers.dispatcher import async_dispatcher_connect import homeassistant.util.dt as dt_util @@ -86,6 +86,7 @@ class UniFiBandwidthSensor(UniFiClient, SensorEntity): DOMAIN = DOMAIN + _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC _attr_native_unit_of_measurement = DATA_MEGABYTES @property @@ -132,6 +133,7 @@ class UniFiUpTimeSensor(UniFiClient, SensorEntity): TYPE = UPTIME_SENSOR _attr_device_class = DEVICE_CLASS_TIMESTAMP + _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC def __init__(self, client, controller): """Set up tracked client.""" diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index eebd5014cb5..03cd9056830 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -15,6 +15,7 @@ from aiounifi.events import ( ) from homeassistant.components.switch import DOMAIN, SwitchEntity +from homeassistant.const import ENTITY_CATEGORY_CONFIG from homeassistant.core import callback from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import DeviceInfo @@ -183,6 +184,8 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchEntity, RestoreEntity): DOMAIN = DOMAIN TYPE = POE_SWITCH + _attr_entity_category = ENTITY_CATEGORY_CONFIG + def __init__(self, client, controller): """Set up POE switch.""" super().__init__(client, controller) @@ -270,6 +273,8 @@ class UniFiBlockClientSwitch(UniFiClient, SwitchEntity): DOMAIN = DOMAIN TYPE = BLOCK_SWITCH + _attr_entity_category = ENTITY_CATEGORY_CONFIG + def __init__(self, client, controller): """Set up block switch.""" super().__init__(client, controller) @@ -319,6 +324,8 @@ class UniFiDPIRestrictionSwitch(UniFiBase, SwitchEntity): DOMAIN = DOMAIN TYPE = DPI_SWITCH + _attr_entity_category = ENTITY_CATEGORY_CONFIG + @property def key(self) -> Any: """Return item key.""" diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index fbf697e295f..2df279090f3 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -15,6 +15,8 @@ from homeassistant.components.unifi.const import ( CONF_TRACK_DEVICES, DOMAIN as UNIFI_DOMAIN, ) +from homeassistant.const import ENTITY_CATEGORY_DIAGNOSTIC +from homeassistant.helpers import entity_registry as er from homeassistant.helpers.dispatcher import async_dispatcher_send import homeassistant.util.dt as dt_util @@ -74,6 +76,12 @@ async def test_bandwidth_sensors(hass, aioclient_mock, mock_unifi_websocket): assert hass.states.get("sensor.wireless_client_rx").state == "2345.0" assert hass.states.get("sensor.wireless_client_tx").state == "6789.0" + ent_reg = er.async_get(hass) + assert ( + ent_reg.async_get("sensor.wired_client_rx").entity_category + == ENTITY_CATEGORY_DIAGNOSTIC + ) + # Verify state update wireless_client["rx_bytes"] = 3456000000 @@ -179,6 +187,12 @@ async def test_uptime_sensors( assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1 assert hass.states.get("sensor.client1_uptime").state == "2021-01-01T01:00:00+00:00" + ent_reg = er.async_get(hass) + assert ( + ent_reg.async_get("sensor.client1_uptime").entity_category + == ENTITY_CATEGORY_DIAGNOSTIC + ) + # Verify normal new event doesn't change uptime # 4 seconds has passed diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index ad277f18a8d..b870465e2bc 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -16,6 +16,7 @@ from homeassistant.components.unifi.const import ( DOMAIN as UNIFI_DOMAIN, ) from homeassistant.components.unifi.switch import POE_SWITCH +from homeassistant.const import ENTITY_CATEGORY_CONFIG from homeassistant.helpers import entity_registry as er from homeassistant.helpers.dispatcher import async_dispatcher_send @@ -373,6 +374,14 @@ async def test_switches(hass, aioclient_mock): assert dpi_switch.state == "on" assert dpi_switch.attributes["icon"] == "mdi:network" + ent_reg = er.async_get(hass) + for entry_id in ( + "switch.poe_client_1", + "switch.block_client_1", + "switch.block_media_streaming", + ): + assert ent_reg.async_get(entry_id).entity_category == ENTITY_CATEGORY_CONFIG + # Block and unblock client aioclient_mock.post(