Add entity category to UniFi sensors and switches (#58484)

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

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

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
 

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(