diff --git a/custom_components/mikrotik_router/binary_sensor.py b/custom_components/mikrotik_router/binary_sensor.py index f5e34f9..57afdd5 100644 --- a/custom_components/mikrotik_router/binary_sensor.py +++ b/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.""" diff --git a/custom_components/mikrotik_router/sensor.py b/custom_components/mikrotik_router/sensor.py index 04da9a8..da898b4 100644 --- a/custom_components/mikrotik_router/sensor.py +++ b/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."""