Added entity category support

custom_components/mikrotik_router/binary_sensor.py

@@ -3,7 +3,7 @@
 import logging
 from typing import Any, Dict, Optional
 
-
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorDeviceClass,
@@ -33,6 +33,7 @@ ATTR_LABEL = "label"
 ATTR_GROUP = "group"
 ATTR_PATH = "data_path"
 ATTR_ATTR = "data_attr"
+ATTR_CTGR = "entity_category"
 
 SENSOR_TYPES = {
     "system_fwupdate": {
@@ -40,6 +41,7 @@ SENSOR_TYPES = {
         ATTR_GROUP: "System",
         ATTR_PATH: "fw-update",
         ATTR_ATTR: "available",
+        ATTR_CTGR: EntityCategory.DIAGNOSTIC,
     },
 }
 
@@ -236,6 +238,11 @@ class MikrotikControllerBinarySensor(BinarySensorEntity):
             self._type = {}
             self._attr = None
 
+        if ATTR_CTGR in self._type:
+            self._entity_category = self._type[ATTR_CTGR]
+        else:
+            self._entity_category = None
+
         self._device_class = None
         self._state = None
         self._icon = None
@@ -262,6 +269,14 @@ class MikrotikControllerBinarySensor(BinarySensorEntity):
         """Return if controller is available."""
         return self._ctrl.connected()
 
+    @property
+    def entity_category(self) -> str:
+        """Return entity category"""
+        if self._entity_category:
+            return self._entity_category
+
+        return None
+
     @property
     def device_info(self) -> Dict[str, Any]:
         """Return a description for device registry."""

custom_components/mikrotik_router/sensor.py

@@ -11,6 +11,7 @@ from homeassistant.const import (
     TEMP_CELSIUS,
 )
 
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.components.sensor import SensorDeviceClass
 
 from .const import (
@@ -51,6 +52,7 @@ ATTR_UNIT_ATTR = "unit_attr"
 ATTR_GROUP = "group"
 ATTR_PATH = "data_path"
 ATTR_ATTR = "data_attr"
+ATTR_CTGR = "entity_category"
 
 SENSOR_TYPES = {
     "system_temperature": {
@@ -115,6 +117,7 @@ SENSOR_TYPES = {
         ATTR_GROUP: "System",
         ATTR_PATH: "resource",
         ATTR_ATTR: "uptime",
+        ATTR_CTGR: EntityCategory.DIAGNOSTIC,
     },
     "system_cpu-load": {
         ATTR_DEVICE_CLASS: None,
@@ -366,6 +369,11 @@ class MikrotikControllerSensor(SensorEntity):
         self._device_class = None
         self._state = None
 
+        if ATTR_CTGR in self._type:
+            self._entity_category = self._type[ATTR_CTGR]
+        else:
+            self._entity_category = None
+
         if ATTR_ICON in self._type:
             self._icon = self._type[ATTR_ICON]
         else:
@@ -401,6 +409,14 @@ class MikrotikControllerSensor(SensorEntity):
 
         return None
 
+    @property
+    def entity_category(self) -> str:
+        """Return entity category"""
+        if self._entity_category:
+            return self._entity_category
+
+        return None
+
     @property
     def device_class(self) -> Optional[str]:
         """Return the device class."""