added support for per entity services

refactored HA entity implementation

custom_components/mikrotik_router/binary_sensor.py

@@ -3,29 +3,20 @@
 import logging
 from typing import Any
 from collections.abc import Mapping
-from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
 )
-from homeassistant.const import (
-    CONF_NAME,
-    CONF_HOST,
-    ATTR_ATTRIBUTION,
-)
-from homeassistant.core import callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from .helper import format_attribute
 from .const import (
-    DOMAIN,
-    ATTRIBUTION,
     CONF_SENSOR_PPP,
     DEFAULT_SENSOR_PPP,
     CONF_SENSOR_PORT_TRACKER,
     DEFAULT_SENSOR_PORT_TRACKER,
 )
+from .model import model_async_setup_entry, MikrotikEntity
 from .binary_sensor_types import (
-    MikrotikBinarySensorEntityDescription,
     SENSOR_TYPES,
+    SENSOR_SERVICES,
     DEVICE_ATTRIBUTES_IFACE_ETHER,
     DEVICE_ATTRIBUTES_IFACE_SFP,
 )
@@ -37,229 +28,47 @@ _LOGGER = logging.getLogger(__name__)
 #   async_setup_entry
 # ---------------------------
 async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up device tracker for Mikrotik Router component."""
+    """Set up entry for component"""
-    inst = config_entry.data[CONF_NAME]
+    dispatcher = {
-    mikrotik_controller = hass.data[DOMAIN][config_entry.entry_id]
+        "MikrotikBinarySensor": MikrotikBinarySensor,
-    sensors = {}
+        "MikrotikPPPSecretBinarySensor": MikrotikPPPSecretBinarySensor,
-
+        "MikrotikPortBinarySensor": MikrotikPortBinarySensor,
-    @callback
+    }
-    def update_controller():
+    await model_async_setup_entry(
-        """Update the values of the controller."""
+        hass,
-        update_items(
-            inst, config_entry, mikrotik_controller, async_add_entities, sensors
-        )
-
-    mikrotik_controller.listeners.append(
-        async_dispatcher_connect(
-            hass, mikrotik_controller.signal_update, update_controller
-        )
-    )
-    update_controller()
-
-
-# ---------------------------
-#   update_items
-# ---------------------------
-@callback
-def update_items(inst, config_entry, mikrotik_controller, async_add_entities, sensors):
-    """Update sensor state from the controller."""
-    new_sensors = []
-
-    for sensor, sid_func in zip(
-        # Sensor type name
-        ["ppp_tracker", "interface"],
-        # Entity function
-        [MikrotikControllerPPPSecretBinarySensor, MikrotikControllerPortBinarySensor],
-    ):
-        if sensor == "interface" and not config_entry.options.get(
-            CONF_SENSOR_PORT_TRACKER, DEFAULT_SENSOR_PORT_TRACKER
-        ):
-            continue
-
-        uid_sensor = SENSOR_TYPES[sensor]
-        for uid in mikrotik_controller.data[uid_sensor.data_path]:
-            uid_data = mikrotik_controller.data[uid_sensor.data_path]
-            if uid_sensor.data_path == "interface" and uid_data[uid]["type"] == "wlan":
-                continue
-
-            item_id = f"{inst}-{sensor}-{uid_data[uid][uid_sensor.data_reference]}"
-            _LOGGER.debug("Updating binary sensor %s", item_id)
-            if item_id in sensors:
-                if sensors[item_id].enabled:
-                    sensors[item_id].async_schedule_update_ha_state()
-                continue
-
-            sensors[item_id] = sid_func(
-                inst=inst,
-                uid=uid,
-                mikrotik_controller=mikrotik_controller,
-                entity_description=uid_sensor,
-                config_entry=config_entry,
-            )
-            new_sensors.append(sensors[item_id])
-
-    for sensor in SENSOR_TYPES:
-        if sensor.startswith("system_"):
-            uid_sensor = SENSOR_TYPES[sensor]
-            item_id = f"{inst}-{sensor}"
-            _LOGGER.debug("Updating binary sensor %s", item_id)
-            if item_id in sensors:
-                if sensors[item_id].enabled:
-                    sensors[item_id].async_schedule_update_ha_state()
-                continue
-
-            sensors[item_id] = MikrotikControllerBinarySensor(
-                inst=inst,
-                uid="",
-                mikrotik_controller=mikrotik_controller,
-                entity_description=uid_sensor,
-                config_entry=config_entry,
-            )
-            new_sensors.append(sensors[item_id])
-
-    if new_sensors:
-        async_add_entities(new_sensors, True)
-
-
-class MikrotikControllerBinarySensor(BinarySensorEntity):
-    """Define an Mikrotik Controller Binary Sensor."""
-
-    def __init__(
-        self,
-        inst,
-        uid: "",
-        mikrotik_controller,
-        entity_description: MikrotikBinarySensorEntityDescription,
         config_entry,
-    ):
+        async_add_entities,
-        """Initialize."""
+        SENSOR_SERVICES,
-        self.entity_description = entity_description
+        SENSOR_TYPES,
-        self._config_entry = config_entry
+        dispatcher,
-        self._inst = inst
+    )
-        self._ctrl = mikrotik_controller
-        self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: ATTRIBUTION}
-        self._uid = uid
-        if self._uid:
-            self._data = mikrotik_controller.data[self.entity_description.data_path][
-                self._uid
-            ]
-        else:
-            self._data = mikrotik_controller.data[self.entity_description.data_path]
 
-    @property
-    def name(self) -> str:
-        """Return the name."""
-        if self._uid:
-            if self.entity_description.name:
-                return f"{self._inst} {self._data[self.entity_description.data_name]} {self.entity_description.name}"
 
-            return f"{self._inst} {self._data[self.entity_description.data_name]}"
+# ---------------------------
-        else:
+#   MikrotikBinarySensor
-            return f"{self._inst} {self.entity_description.name}"
+# ---------------------------
-
+class MikrotikBinarySensor(MikrotikEntity, BinarySensorEntity):
-    @property
+    """Define an Mikrotik Controller Binary Sensor."""
-    def unique_id(self) -> str:
-        """Return a unique id for this entity."""
-        if self._uid:
-            return f"{self._inst.lower()}-{self.entity_description.key}-{self._data[self.entity_description.data_reference].lower()}"
-        else:
-            return f"{self._inst.lower()}-{self.entity_description.key}"
-
-    @property
-    def available(self) -> bool:
-        """Return if controller is available."""
-        return self._ctrl.connected()
 
     @property
     def is_on(self) -> bool:
         """Return true if device is on."""
-        return self._data[self.entity_description.data_is_on]
+        return self._data[self.entity_description.data_attribute]
 
     @property
     def icon(self) -> str:
         """Return the icon."""
         if self.entity_description.icon_enabled:
-            if self._data[self.entity_description.data_is_on]:
+            if self._data[self.entity_description.data_attribute]:
                 return self.entity_description.icon_enabled
             else:
                 return self.entity_description.icon_disabled
 
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return a description for device registry."""
-        dev_connection = DOMAIN
-        dev_connection_value = self.entity_description.data_reference
-        dev_group = self.entity_description.ha_group
-        if self.entity_description.ha_group == "System":
-            dev_group = self._ctrl.data["resource"]["board-name"]
-            dev_connection_value = self._ctrl.data["routerboard"]["serial-number"]
-
-        if self.entity_description.ha_group.startswith("data__"):
-            dev_group = self.entity_description.ha_group[6:]
-            if dev_group in self._data:
-                dev_group = self._data[dev_group]
-                dev_connection_value = dev_group
-
-        if self.entity_description.ha_connection:
-            dev_connection = self.entity_description.ha_connection
-
-        if self.entity_description.ha_connection_value:
-            dev_connection_value = self.entity_description.ha_connection_value
-            if dev_connection_value.startswith("data__"):
-                dev_connection_value = dev_connection_value[6:]
-                dev_connection_value = self._data[dev_connection_value]
-
-        info = DeviceInfo(
-            connections={(dev_connection, f"{dev_connection_value}")},
-            identifiers={(dev_connection, f"{dev_connection_value}")},
-            default_name=f"{self._inst} {dev_group}",
-            default_model=f"{self._ctrl.data['resource']['board-name']}",
-            default_manufacturer=f"{self._ctrl.data['resource']['platform']}",
-            sw_version=f"{self._ctrl.data['resource']['version']}",
-            configuration_url=f"http://{self._ctrl.config_entry.data[CONF_HOST]}",
-            via_device=(DOMAIN, f"{self._ctrl.data['routerboard']['serial-number']}"),
-        )
-
-        if "mac-address" in self.entity_description.data_reference:
-            dev_group = self._data[self.entity_description.data_name]
-            dev_manufacturer = ""
-            if dev_connection_value in self._ctrl.data["host"]:
-                dev_group = self._ctrl.data["host"][dev_connection_value]["host-name"]
-                dev_manufacturer = self._ctrl.data["host"][dev_connection_value][
-                    "manufacturer"
-                ]
-
-            info = DeviceInfo(
-                connections={(dev_connection, f"{dev_connection_value}")},
-                default_name=f"{dev_group}",
-                deafult_manufacturer=f"{dev_manufacturer}",
-                via_device=(
-                    DOMAIN,
-                    f"{self._ctrl.data['routerboard']['serial-number']}",
-                ),
-            )
-
-        return info
-
-    @property
-    def extra_state_attributes(self) -> Mapping[str, Any]:
-        """Return the state attributes."""
-        attributes = super().extra_state_attributes
-        for variable in self.entity_description.data_attributes_list:
-            if variable in self._data:
-                attributes[format_attribute(variable)] = self._data[variable]
-
-        return attributes
-
-    async def async_added_to_hass(self):
-        """Run when entity about to be added to hass."""
-        _LOGGER.debug("New binary sensor %s (%s)", self._inst, self.unique_id)
-
 
 # ---------------------------
-#   MikrotikControllerPPPSecretBinarySensor
+#   MikrotikPPPSecretBinarySensor
 # ---------------------------
-class MikrotikControllerPPPSecretBinarySensor(MikrotikControllerBinarySensor):
+class MikrotikPPPSecretBinarySensor(MikrotikBinarySensor):
     """Representation of a network device."""
 
     @property
@@ -270,24 +79,22 @@ class MikrotikControllerPPPSecretBinarySensor(MikrotikControllerBinarySensor):
     @property
     def is_on(self) -> bool:
         """Return true if device is on."""
-        if not self.option_sensor_ppp:
+        return (
-            return False
+            self._data[self.entity_description.data_attribute]
-
+            if self.option_sensor_ppp
-        return self._data[self.entity_description.data_is_on]
+            else False
+        )
 
     @property
     def available(self) -> bool:
         """Return if controller is available."""
-        if not self.option_sensor_ppp:
+        return self._ctrl.connected() if self.option_sensor_ppp else False
-            return False
-
-        return self._ctrl.connected()
 
 
 # ---------------------------
-#   MikrotikControllerPortBinarySensor
+#   MikrotikPortBinarySensor
 # ---------------------------
-class MikrotikControllerPortBinarySensor(MikrotikControllerBinarySensor):
+class MikrotikPortBinarySensor(MikrotikBinarySensor):
     """Representation of a network port."""
 
     @property
@@ -300,15 +107,12 @@ class MikrotikControllerPortBinarySensor(MikrotikControllerBinarySensor):
     @property
     def available(self) -> bool:
         """Return if controller is available."""
-        if not self.option_sensor_port_tracker:
+        return self._ctrl.connected() if self.option_sensor_port_tracker else False
-            return False
-
-        return self._ctrl.connected()
 
     @property
     def icon(self) -> str:
         """Return the icon."""
-        if self._data[self.entity_description.data_is_on]:
+        if self._data[self.entity_description.data_attribute]:
             icon = self.entity_description.icon_enabled
         else:
             icon = self.entity_description.icon_disabled

custom_components/mikrotik_router/binary_sensor_types.py

@@ -77,11 +77,13 @@ class MikrotikBinarySensorEntityDescription(BinarySensorEntityDescription):
     ha_connection: str = ""
     ha_connection_value: str = ""
     data_path: str = ""
-    data_is_on: str = "available"
+    data_attribute: str = "available"
     data_name: str = ""
+    data_name_comment: bool = False
     data_uid: str = ""
     data_reference: str = ""
     data_attributes_list: List = field(default_factory=lambda: [])
+    func: str = "MikrotikBinarySensor"
 
 
 SENSOR_TYPES = {
@@ -108,11 +110,12 @@ SENSOR_TYPES = {
         ha_connection=DOMAIN,
         ha_connection_value="PPP",
         data_path="ppp_secret",
-        data_is_on="connected",
+        data_attribute="connected",
         data_name="name",
         data_uid="name",
         data_reference="name",
         data_attributes_list=DEVICE_ATTRIBUTES_PPP_SECRET,
+        func="MikrotikPPPSecretBinarySensor",
     ),
     "interface": MikrotikBinarySensorEntityDescription(
         key="interface",
@@ -124,10 +127,13 @@ SENSOR_TYPES = {
         ha_connection=CONNECTION_NETWORK_MAC,
         ha_connection_value="data__port-mac-address",
         data_path="interface",
-        data_is_on="running",
+        data_attribute="running",
         data_name="name",
         data_uid="default-name",
         data_reference="default-name",
         data_attributes_list=DEVICE_ATTRIBUTES_IFACE,
+        func="MikrotikPortBinarySensor",
     ),
 }
+
+SENSOR_SERVICES = {}

custom_components/mikrotik_router/button.py

@@ -1,198 +1,54 @@
 """Support for the Mikrotik Router buttons."""
 
 import logging
-from typing import Any, Dict
 from homeassistant.components.button import ButtonEntity
-from homeassistant.const import CONF_NAME, ATTR_ATTRIBUTION
+from .model import model_async_setup_entry, MikrotikEntity
-from homeassistant.core import callback
+from .button_types import (
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
+    SENSOR_TYPES,
-from homeassistant.helpers.restore_state import RestoreEntity
+    SENSOR_SERVICES,
-from .helper import format_attribute
+)
-from .const import DOMAIN, ATTRIBUTION
 
 _LOGGER = logging.getLogger(__name__)
 
-DEVICE_ATTRIBUTES_SCRIPT = [
-    "last-started",
-    "run-count",
-]
-
 
 # ---------------------------
 #   async_setup_entry
 # ---------------------------
 async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up buttons for Mikrotik Router component."""
+    """Set up entry for component"""
-    inst = config_entry.data[CONF_NAME]
+    dispatcher = {
-    mikrotik_controller = hass.data[DOMAIN][config_entry.entry_id]
+        "MikrotikButton": MikrotikButton,
-    buttons = {}
+        "MikrotikScriptButton": MikrotikScriptButton,
-
+    }
-    @callback
+    await model_async_setup_entry(
-    def update_controller():
+        hass,
-        """Update the values of the controller."""
+        config_entry,
-        update_items(inst, mikrotik_controller, async_add_entities, buttons)
+        async_add_entities,
-
+        SENSOR_SERVICES,
-    mikrotik_controller.listeners.append(
+        SENSOR_TYPES,
-        async_dispatcher_connect(
+        dispatcher,
-            hass, mikrotik_controller.signal_update, update_controller
-        )
     )
 
-    update_controller()
-
 
 # ---------------------------
-#   update_items
+#   MikrotikButton
 # ---------------------------
-@callback
+class MikrotikButton(MikrotikEntity, ButtonEntity):
-def update_items(inst, mikrotik_controller, async_add_entities, buttons):
-    """Update device button state from the controller."""
-    new_buttons = []
-
-    # Add buttons
-    for sid, sid_uid, sid_name, sid_ref, sid_attr, sid_func in zip(
-        # Data point name
-        [
-            "script",
-        ],
-        # Data point unique id
-        [
-            "name",
-        ],
-        # Entry Name
-        [
-            "name",
-        ],
-        # Entry Unique id
-        [
-            "name",
-        ],
-        # Attr
-        [
-            DEVICE_ATTRIBUTES_SCRIPT,
-        ],
-        # Button function
-        [
-            MikrotikControllerScriptButton,
-        ],
-    ):
-        for uid in mikrotik_controller.data[sid]:
-            item_id = f"{inst}-{sid}-{mikrotik_controller.data[sid][uid][sid_uid]}"
-
-            _LOGGER.debug("Updating button %s", item_id)
-            if item_id in buttons:
-                if buttons[item_id].enabled:
-                    buttons[item_id].async_schedule_update_ha_state()
-                continue
-
-            # Create new entity
-            sid_data = {
-                "sid": sid,
-                "sid_uid": sid_uid,
-                "sid_name": sid_name,
-                "sid_ref": sid_ref,
-                "sid_attr": sid_attr,
-            }
-            buttons[item_id] = sid_func(inst, uid, mikrotik_controller, sid_data)
-            new_buttons.append(buttons[item_id])
-
-    if new_buttons:
-        async_add_entities(new_buttons)
-
-
-# ---------------------------
-#   MikrotikControllerButton
-# ---------------------------
-class MikrotikControllerButton(ButtonEntity, RestoreEntity):
     """Representation of a button."""
 
-    def __init__(self, inst, uid, mikrotik_controller, sid_data):
-        """Initialize."""
-        self._sid_data = sid_data
-        self._inst = inst
-        self._ctrl = mikrotik_controller
-        self._data = mikrotik_controller.data[self._sid_data["sid"]][uid]
-
-        self._attrs = {
-            ATTR_ATTRIBUTION: ATTRIBUTION,
-        }
-
-    async def async_added_to_hass(self):
-        """Run when entity about to be added to hass."""
-        _LOGGER.debug(
-            "New button %s (%s %s)",
-            self._inst,
-            self._sid_data["sid"],
-            self._data[self._sid_data["sid_uid"]],
-        )
-
     async def async_update(self):
         """Synchronize state with controller."""
 
-    @property
-    def available(self) -> bool:
-        """Return if controller is available."""
-        return self._ctrl.connected()
-
-    @property
-    def name(self) -> str:
-        """Return the name."""
-        return f"{self._inst} {self._sid_data['sid']} {self._data[self._sid_data['sid_name']]}"
-
-    @property
-    def unique_id(self) -> str:
-        """Return a unique id for this entity."""
-        return f"{self._inst.lower()}-{self._sid_data['sid']}_button-{self._data[self._sid_data['sid_ref']]}"
-
-    @property
-    def extra_state_attributes(self) -> Dict[str, Any]:
-        """Return the state attributes."""
-        attributes = self._attrs
-
-        for variable in self._sid_data["sid_attr"]:
-            if variable in self._data:
-                attributes[format_attribute(variable)] = self._data[variable]
-
-        return attributes
-
     async def async_press(self) -> None:
         pass
 
 
 # ---------------------------
-#   MikrotikControllerScriptButton
+#   MikrotikScriptButton
 # ---------------------------
-class MikrotikControllerScriptButton(MikrotikControllerButton):
+class MikrotikScriptButton(MikrotikButton):
     """Representation of a script button."""
 
-    def __init__(self, inst, uid, mikrotik_controller, sid_data):
-        """Initialize."""
-        super().__init__(inst, uid, mikrotik_controller, sid_data)
-
-    @property
-    def icon(self) -> str:
-        """Return the icon."""
-        return "mdi:script-text-outline"
-
-    @property
-    def device_info(self) -> Dict[str, Any]:
-        """Return a description for device registry."""
-        info = {
-            "identifiers": {
-                (
-                    DOMAIN,
-                    "serial-number",
-                    f"{self._ctrl.data['routerboard']['serial-number']}",
-                    "button",
-                    "Scripts",
-                )
-            },
-            "manufacturer": self._ctrl.data["resource"]["platform"],
-            "model": self._ctrl.data["resource"]["board-name"],
-            "name": f"{self._inst} Scripts",
-        }
-        return info
-
     async def async_press(self) -> None:
         """Process the button press."""
         self._ctrl.run_script(self._data["name"])

custom_components/mikrotik_router/button_types.py

@@ -0,0 +1,52 @@
+"""Definitions for Mikrotik Router button entities."""
+from dataclasses import dataclass, field
+from typing import List
+from homeassistant.components.sensor import (
+    SensorEntityDescription,
+)
+from .const import DOMAIN
+
+
+DEVICE_ATTRIBUTES_SCRIPT = [
+    "last-started",
+    "run-count",
+]
+
+
+@dataclass
+class MikrotikButtonEntityDescription(SensorEntityDescription):
+    """Class describing mikrotik entities."""
+
+    ha_group: str = ""
+    ha_connection: str = ""
+    ha_connection_value: str = ""
+    data_path: str = ""
+    data_attribute: str = ""
+    data_name: str = ""
+    data_name_comment: bool = False
+    data_uid: str = ""
+    data_reference: str = ""
+    data_attributes_list: List = field(default_factory=lambda: [])
+    func: str = "MikrotikButton"
+
+
+SENSOR_TYPES = {
+    "script": MikrotikButtonEntityDescription(
+        key="script",
+        name="script",
+        icon="mdi:script-text-outline",
+        device_class=None,
+        entity_category=None,
+        ha_group="Script",
+        ha_connection=DOMAIN,
+        ha_connection_value="Script",
+        data_path="script",
+        data_name="name",
+        data_uid="name",
+        data_reference="name",
+        data_attributes_list=DEVICE_ATTRIBUTES_SCRIPT,
+        func="MikrotikScriptButton",
+    ),
+}
+
+SENSOR_SERVICES = {}

custom_components/mikrotik_router/device_tracker.py

@@ -1,35 +1,21 @@
 """Support for the Mikrotik Router device tracker."""
-
 import logging
-from typing import Any, Dict
+from typing import Any
 from collections.abc import Mapping
 from datetime import timedelta
 from homeassistant.components.device_tracker.config_entry import ScannerEntity
 from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER
-from homeassistant.const import (
+from homeassistant.const import STATE_NOT_HOME
-    CONF_NAME,
-    CONF_HOST,
-    ATTR_ATTRIBUTION,
-    STATE_NOT_HOME,
-)
-from homeassistant.core import callback
-from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
-from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.util.dt import get_age, utcnow
-from .helper import format_attribute, format_value
+from .helper import format_attribute
 from .const import (
-    DOMAIN,
-    ATTRIBUTION,
     CONF_TRACK_HOSTS,
     DEFAULT_TRACK_HOSTS,
     CONF_TRACK_HOSTS_TIMEOUT,
     DEFAULT_TRACK_HOST_TIMEOUT,
 )
-from .device_tracker_types import (
+from .model import model_async_setup_entry, MikrotikEntity
-    MikrotikDeviceTrackerEntityDescription,
+from .device_tracker_types import SENSOR_TYPES, SENSOR_SERVICES
-    SENSOR_TYPES,
-)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -38,114 +24,31 @@ _LOGGER = logging.getLogger(__name__)
 #   async_setup_entry
 # ---------------------------
 async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up device tracker for Mikrotik Router component."""
+    """Set up entry for component"""
-    inst = config_entry.data[CONF_NAME]
+    dispatcher = {
-    mikrotik_controller = hass.data[DOMAIN][config_entry.entry_id]
+        "MikrotikDeviceTracker": MikrotikDeviceTracker,
-    trackers = {}
+        "MikrotikHostDeviceTracker": MikrotikHostDeviceTracker,
-
+    }
-    @callback
+    await model_async_setup_entry(
-    def update_controller():
+        hass,
-        """Update the values of the controller."""
+        config_entry,
-        update_items(
+        async_add_entities,
-            inst, config_entry, mikrotik_controller, async_add_entities, trackers
+        SENSOR_SERVICES,
-        )
+        SENSOR_TYPES,
-
+        dispatcher,
-    mikrotik_controller.listeners.append(
-        async_dispatcher_connect(
-            hass, mikrotik_controller.signal_update, update_controller
-        )
     )
 
-    update_controller()
-
 
 # ---------------------------
-#   update_items
+#   MikrotikDeviceTracker
 # ---------------------------
-@callback
+class MikrotikDeviceTracker(MikrotikEntity, ScannerEntity):
-def update_items(inst, config_entry, mikrotik_controller, async_add_entities, trackers):
-    """Update trackers device state from the controller."""
-    new_trackers = []
-
-    for sensor, sid_func in zip(
-        # Sensor type name
-        ["host"],
-        # Entity function
-        [MikrotikControllerHostDeviceTracker],
-    ):
-        uid_sensor = SENSOR_TYPES[sensor]
-        if (
-            # Skip if host tracking is disabled
-            sensor == "host"
-            and not config_entry.options.get(CONF_TRACK_HOSTS, DEFAULT_TRACK_HOSTS)
-        ):
-            continue
-
-        for uid in mikrotik_controller.data[uid_sensor.data_path]:
-            uid_data = mikrotik_controller.data[uid_sensor.data_path]
-            item_id = f"{inst}-{sensor}-{uid_data[uid][uid_sensor.data_reference]}"
-            _LOGGER.debug("Updating device tracker %s", item_id)
-            if item_id in trackers:
-                if trackers[item_id].enabled:
-                    trackers[item_id].async_schedule_update_ha_state()
-                continue
-
-            trackers[item_id] = sid_func(
-                inst=inst,
-                uid=uid,
-                mikrotik_controller=mikrotik_controller,
-                entity_description=uid_sensor,
-                config_entry=config_entry,
-            )
-            new_trackers.append(trackers[item_id])
-
-    # Register new entities
-    if new_trackers:
-        async_add_entities(new_trackers)
-
-
-# ---------------------------
-#   MikrotikControllerDeviceTracker
-# ---------------------------
-class MikrotikControllerDeviceTracker(ScannerEntity):
     """Representation of a device tracker."""
 
-    def __init__(
-        self,
-        inst,
-        uid: "",
-        mikrotik_controller,
-        entity_description: MikrotikDeviceTrackerEntityDescription,
-        config_entry,
-    ):
-        """Initialize."""
-        self.entity_description = entity_description
-        self._config_entry = config_entry
-        self._inst = inst
-        self._ctrl = mikrotik_controller
-        self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: ATTRIBUTION}
-        self._data = mikrotik_controller.data[self.entity_description.data_path][uid]
-
-    @property
-    def name(self) -> str:
-        """Return the name."""
-        if self.entity_description.name:
-            return f"{self._inst} {self._data[self.entity_description.data_name]} {self.entity_description.name}"
-
-        return f"{self._inst} {self._data[self.entity_description.data_name]}"
-
-    @property
-    def unique_id(self) -> str:
-        """Return a unique id for this entity."""
-        return f"{self._inst.lower()}-{self.entity_description.key}-{self._data[self.entity_description.data_reference].lower()}"
-
     @property
     def ip_address(self) -> str:
         """Return the primary ip address of the device."""
-        if "address" in self._data:
+        return self._data["address"] if "address" in self._data else None
-            return self._data["address"]
-
-        return None
 
     @property
     def mac_address(self) -> str:
@@ -153,7 +56,7 @@ class MikrotikControllerDeviceTracker(ScannerEntity):
         if self.entity_description.data_reference in self._data:
             return self._data[self.entity_description.data_reference]
 
-        return None
+        return ""
 
     @property
     def hostname(self) -> str:
@@ -161,90 +64,23 @@ class MikrotikControllerDeviceTracker(ScannerEntity):
         if self.entity_description.data_name in self._data:
             return self._data[self.entity_description.data_name]
 
-        return None
+        return ""
 
     @property
     def is_connected(self) -> bool:
         """Return true if device is connected."""
-        return self._data[self.entity_description.data_is_on]
+        return self._data[self.entity_description.data_attribute]
-
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return a description for device registry."""
-        dev_connection = DOMAIN
-        dev_connection_value = self.entity_description.data_reference
-        dev_group = self.entity_description.ha_group
-        if self.entity_description.ha_group.startswith("data__"):
-            dev_group = self.entity_description.ha_group[6:]
-            if dev_group in self._data:
-                dev_group = self._data[dev_group]
-                dev_connection_value = dev_group
-
-        if self.entity_description.ha_connection:
-            dev_connection = self.entity_description.ha_connection
-
-        if self.entity_description.ha_connection_value:
-            dev_connection_value = self.entity_description.ha_connection_value
-            if dev_connection_value.startswith("data__"):
-                dev_connection_value = dev_connection_value[6:]
-                dev_connection_value = self._data[dev_connection_value]
-
-        info = DeviceInfo(
-            connections={(dev_connection, f"{dev_connection_value}")},
-            identifiers={(dev_connection, f"{dev_connection_value}")},
-            default_name=f"{self._inst} {dev_group}",
-            default_model=f"{self._ctrl.data['resource']['board-name']}",
-            default_manufacturer=f"{self._ctrl.data['resource']['platform']}",
-            sw_version=f"{self._ctrl.data['resource']['version']}",
-            configuration_url=f"http://{self._ctrl.config_entry.data[CONF_HOST]}",
-            via_device=(DOMAIN, f"{self._ctrl.data['routerboard']['serial-number']}"),
-        )
-
-        if "mac-address" in self.entity_description.data_reference:
-            dev_group = self._data[self.entity_description.data_name]
-            dev_manufacturer = ""
-            if dev_connection_value in self._ctrl.data["host"]:
-                dev_group = self._ctrl.data["host"][dev_connection_value]["host-name"]
-                dev_manufacturer = self._ctrl.data["host"][dev_connection_value][
-                    "manufacturer"
-                ]
-
-            info = DeviceInfo(
-                connections={(dev_connection, f"{dev_connection_value}")},
-                default_name=f"{dev_group}",
-                default_manufacturer=f"{dev_manufacturer}",
-                via_device=(
-                    DOMAIN,
-                    f"{self._ctrl.data['routerboard']['serial-number']}",
-                ),
-            )
-
-        return info
-
-    @property
-    def extra_state_attributes(self) -> Mapping[str, Any]:
-        """Return the state attributes."""
-        attributes = super().extra_state_attributes
-        for variable in self.entity_description.data_attributes_list:
-            if variable in self._data:
-                attributes[format_attribute(variable)] = self._data[variable]
-
-        return attributes
 
     @property
     def source_type(self) -> str:
         """Return the source type of the port."""
         return SOURCE_TYPE_ROUTER
 
-    async def async_added_to_hass(self):
-        """Run when entity about to be added to hass."""
-        _LOGGER.debug("New device tracker %s (%s)", self._inst, self.unique_id)
-
 
 # ---------------------------
-#   MikrotikControllerHostDeviceTracker
+#   MikrotikHostDeviceTracker
 # ---------------------------
-class MikrotikControllerHostDeviceTracker(MikrotikControllerDeviceTracker):
+class MikrotikHostDeviceTracker(MikrotikDeviceTracker):
     """Representation of a network device."""
 
     @property
@@ -277,7 +113,7 @@ class MikrotikControllerHostDeviceTracker(MikrotikControllerDeviceTracker):
             return False
 
         if self._data["source"] in ["capsman", "wireless"]:
-            return self._data[self.entity_description.data_is_on]
+            return self._data[self.entity_description.data_attribute]
 
         if (
             self._data["last-seen"]
@@ -292,7 +128,7 @@ class MikrotikControllerHostDeviceTracker(MikrotikControllerDeviceTracker):
     def icon(self) -> str:
         """Return the icon."""
         if self._data["source"] in ["capsman", "wireless"]:
-            if self._data[self.entity_description.data_is_on]:
+            if self._data[self.entity_description.data_attribute]:
                 return self.entity_description.icon_enabled
             else:
                 return self.entity_description.icon_disabled
@@ -308,12 +144,10 @@ class MikrotikControllerHostDeviceTracker(MikrotikControllerDeviceTracker):
     @property
     def state(self) -> str:
         """Return the state of the device."""
-        if self.is_connected:
+        return self._ctrl.option_zone if self.is_connected else STATE_NOT_HOME
-            return self._ctrl.option_zone
-        return STATE_NOT_HOME
 
     @property
-    def extra_state_attributes(self) -> Dict[str, Any]:
+    def extra_state_attributes(self) -> Mapping[str, Any]:
         """Return the state attributes."""
         attributes = super().extra_state_attributes
         if self._data["last-seen"]:

custom_components/mikrotik_router/device_tracker_types.py

@@ -26,11 +26,12 @@ class MikrotikDeviceTrackerEntityDescription(SwitchEntityDescription):
     ha_connection: str = ""
     ha_connection_value: str = ""
     data_path: str = ""
-    data_is_on: str = "available"
+    data_attribute: str = "available"
     data_name: str = ""
     data_uid: str = ""
     data_reference: str = ""
     data_attributes_list: List = field(default_factory=lambda: [])
+    func: str = "MikrotikDeviceTracker"
 
 
 SENSOR_TYPES = {
@@ -47,5 +48,8 @@ SENSOR_TYPES = {
         data_uid="mac-address",
         data_reference="mac-address",
         data_attributes_list=DEVICE_ATTRIBUTES_HOST,
+        func="MikrotikHostDeviceTracker",
     ),
 }
+
+SENSOR_SERVICES = {}

custom_components/mikrotik_router/model.py

@@ -0,0 +1,296 @@
+"""Mikrotik HA shared entity model"""
+from logging import getLogger
+from typing import Any
+from collections.abc import Mapping
+from homeassistant.helpers import entity_platform
+from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.core import callback
+from homeassistant.const import ATTR_ATTRIBUTION, CONF_NAME, CONF_HOST
+from .helper import format_attribute
+from .const import (
+    DOMAIN,
+    ATTRIBUTION,
+    CONF_SENSOR_PORT_TRAFFIC,
+    DEFAULT_SENSOR_PORT_TRAFFIC,
+    CONF_TRACK_HOSTS,
+    DEFAULT_TRACK_HOSTS,
+)
+
+_LOGGER = getLogger(__name__)
+
+
+# ---------------------------
+#   model_async_setup_entry
+# ---------------------------
+async def model_async_setup_entry(
+    hass, config_entry, async_add_entities, sensor_services, sensor_types, dispatcher
+):
+    inst = config_entry.data[CONF_NAME]
+    mikrotik_controller = hass.data[DOMAIN][config_entry.entry_id]
+    sensors = {}
+
+    platform = entity_platform.async_get_current_platform()
+    for service in sensor_services:
+        platform.async_register_entity_service(service[0], service[1], service[2])
+
+    @callback
+    def update_controller():
+        """Update the values of the controller"""
+        model_update_items(
+            inst,
+            config_entry,
+            mikrotik_controller,
+            async_add_entities,
+            sensors,
+            dispatcher,
+            sensor_types,
+        )
+
+    mikrotik_controller.listeners.append(
+        async_dispatcher_connect(
+            hass, mikrotik_controller.signal_update, update_controller
+        )
+    )
+    update_controller()
+
+
+# ---------------------------
+#   model_update_items
+# ---------------------------
+def model_update_items(
+    inst,
+    config_entry,
+    mikrotik_controller,
+    async_add_entities,
+    sensors,
+    dispatcher,
+    sensor_types,
+):
+    def _register_entity(_sensors, _item_id, _uid, _uid_sensor):
+        _LOGGER.debug("Updating entity %s", _item_id)
+        if _item_id in _sensors:
+            if _sensors[_item_id].enabled:
+                _sensors[_item_id].async_schedule_update_ha_state()
+            return None
+
+        return dispatcher[_uid_sensor.func](
+            inst=inst,
+            uid=_uid,
+            mikrotik_controller=mikrotik_controller,
+            entity_description=_uid_sensor,
+        )
+
+    new_sensors = []
+    for sensor in sensor_types:
+        uid_sensor = sensor_types[sensor]
+        if not uid_sensor.data_reference:
+            uid_sensor = sensor_types[sensor]
+            if (
+                uid_sensor.data_attribute
+                not in mikrotik_controller.data[uid_sensor.data_path]
+                or mikrotik_controller.data[uid_sensor.data_path][
+                    uid_sensor.data_attribute
+                ]
+                == "unknown"
+            ):
+                continue
+
+            item_id = f"{inst}-{sensor}"
+            _LOGGER.debug("Updating entity %s", item_id)
+            if tmp := _register_entity(sensors, item_id, "", uid_sensor):
+                sensors[item_id] = tmp
+                new_sensors.append(sensors[item_id])
+        else:
+            # Sensors
+            if sensor.startswith("traffic_") and not config_entry.options.get(
+                CONF_SENSOR_PORT_TRAFFIC, DEFAULT_SENSOR_PORT_TRAFFIC
+            ):
+                continue
+
+            for uid in mikrotik_controller.data[uid_sensor.data_path]:
+                uid_data = mikrotik_controller.data[uid_sensor.data_path]
+
+                # Sensors
+                if (
+                    uid_sensor.data_path == "interface"
+                    and uid_data[uid]["type"] == "bridge"
+                ):
+                    continue
+
+                if (
+                    uid_sensor.data_path == "client_traffic"
+                    and uid_sensor.data_attribute not in uid_data[uid].keys()
+                ):
+                    continue
+
+                # Binary sensors
+                if (
+                    uid_sensor.data_path == "interface"
+                    and uid_data[uid]["type"] == "wlan"
+                ):
+                    continue
+
+                # Device Tracker
+                if (
+                    # Skip if host tracking is disabled
+                    sensor == "host"
+                    and not config_entry.options.get(
+                        CONF_TRACK_HOSTS, DEFAULT_TRACK_HOSTS
+                    )
+                ):
+                    continue
+
+                item_id = f"{inst}-{sensor}-{str(uid_data[uid][uid_sensor.data_reference]).lower()}"
+                _LOGGER.debug("Updating entity %s", item_id)
+                if tmp := _register_entity(sensors, item_id, uid, uid_sensor):
+                    sensors[item_id] = tmp
+                    new_sensors.append(sensors[item_id])
+
+    if new_sensors:
+        async_add_entities(new_sensors, True)
+
+
+# ---------------------------
+#   MikrotikEntity
+# ---------------------------
+class MikrotikEntity:
+    """Define entity"""
+
+    def __init__(
+        self,
+        inst,
+        uid: "",
+        mikrotik_controller,
+        entity_description,
+    ):
+        """Initialize entity"""
+        self.entity_description = entity_description
+        self._inst = inst
+        self._ctrl = mikrotik_controller
+        self._config_entry = self._ctrl.config_entry
+        self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: ATTRIBUTION}
+        self._uid = uid
+        if self._uid:
+            self._data = mikrotik_controller.data[self.entity_description.data_path][
+                self._uid
+            ]
+        else:
+            self._data = mikrotik_controller.data[self.entity_description.data_path]
+
+    @property
+    def name(self) -> str:
+        """Return the name for this entity"""
+        if not self._uid:
+            if self.entity_description.data_name_comment and self._data["comment"]:
+                return f"{self._inst} {self._data['comment']}"
+
+            return f"{self._inst} {self.entity_description.name}"
+
+        if self.entity_description.name:
+            if self.entity_description.data_name_comment and self._data["comment"]:
+                return f"{self._inst} {self.entity_description.name} {self._data['comment']}"
+
+            return f"{self._inst} {self._data[self.entity_description.data_name]} {self.entity_description.name}"
+
+        return f"{self._inst} {self._data[self.entity_description.data_name]}"
+
+    @property
+    def unique_id(self) -> str:
+        """Return a unique id for this entity"""
+        if self._uid:
+            return f"{self._inst.lower()}-{self.entity_description.key}-{str(self._data[self.entity_description.data_reference]).lower()}"
+        else:
+            return f"{self._inst.lower()}-{self.entity_description.key}"
+
+    @property
+    def available(self) -> bool:
+        """Return if controller is available"""
+        return self._ctrl.connected()
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return a description for device registry."""
+        dev_connection = DOMAIN
+        dev_connection_value = self.entity_description.data_reference
+        dev_group = self.entity_description.ha_group
+        if self.entity_description.ha_group == "System":
+            dev_group = self._ctrl.data["resource"]["board-name"]
+            dev_connection_value = self._ctrl.data["routerboard"]["serial-number"]
+
+        if self.entity_description.ha_group.startswith("data__"):
+            dev_group = self.entity_description.ha_group[6:]
+            if dev_group in self._data:
+                dev_group = self._data[dev_group]
+                dev_connection_value = dev_group
+
+        if self.entity_description.ha_connection:
+            dev_connection = self.entity_description.ha_connection
+
+        if self.entity_description.ha_connection_value:
+            dev_connection_value = self.entity_description.ha_connection_value
+            if dev_connection_value.startswith("data__"):
+                dev_connection_value = dev_connection_value[6:]
+                dev_connection_value = self._data[dev_connection_value]
+
+        info = DeviceInfo(
+            connections={(dev_connection, f"{dev_connection_value}")},
+            identifiers={(dev_connection, f"{dev_connection_value}")},
+            default_name=f"{self._inst} {dev_group}",
+            default_model=f"{self._ctrl.data['resource']['board-name']}",
+            default_manufacturer=f"{self._ctrl.data['resource']['platform']}",
+            sw_version=f"{self._ctrl.data['resource']['version']}",
+            configuration_url=f"http://{self._ctrl.config_entry.data[CONF_HOST]}",
+            via_device=(DOMAIN, f"{self._ctrl.data['routerboard']['serial-number']}"),
+        )
+
+        if "mac-address" in self.entity_description.data_reference:
+            dev_group = self._data[self.entity_description.data_name]
+            dev_manufacturer = ""
+            if dev_connection_value in self._ctrl.data["host"]:
+                dev_group = self._ctrl.data["host"][dev_connection_value]["host-name"]
+                dev_manufacturer = self._ctrl.data["host"][dev_connection_value][
+                    "manufacturer"
+                ]
+
+            info = DeviceInfo(
+                connections={(dev_connection, f"{dev_connection_value}")},
+                default_name=f"{dev_group}",
+                default_manufacturer=f"{dev_manufacturer}",
+                via_device=(
+                    DOMAIN,
+                    f"{self._ctrl.data['routerboard']['serial-number']}",
+                ),
+            )
+
+        return info
+
+    @property
+    def extra_state_attributes(self) -> Mapping[str, Any]:
+        """Return the state attributes."""
+        attributes = super().extra_state_attributes
+        for variable in self.entity_description.data_attributes_list:
+            if variable in self._data:
+                attributes[format_attribute(variable)] = self._data[variable]
+
+        return attributes
+
+    async def async_added_to_hass(self):
+        """Run when entity about to be added to hass"""
+        _LOGGER.debug("New entity %s (%s)", self._inst, self.unique_id)
+
+    async def start(self):
+        """Dummy run function"""
+        _LOGGER.error("Start functionality does not exist for %s", self.entity_id)
+
+    async def stop(self):
+        """Dummy stop function"""
+        _LOGGER.error("Stop functionality does not exist for %s", self.entity_id)
+
+    async def restart(self):
+        """Dummy restart function"""
+        _LOGGER.error("Restart functionality does not exist for %s", self.entity_id)
+
+    async def reload(self):
+        """Dummy reload function"""
+        _LOGGER.error("Reload functionality does not exist for %s", self.entity_id)

custom_components/mikrotik_router/sensor.py

@@ -3,25 +3,12 @@
 import logging
 from typing import Any, Optional
 from collections.abc import Mapping
-from homeassistant.const import (
-    CONF_NAME,
-    CONF_HOST,
-    ATTR_ATTRIBUTION,
-)
-from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.components.sensor import SensorEntity
-from homeassistant.core import callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from .helper import format_attribute
-from .const import (
+from .model import model_async_setup_entry, MikrotikEntity
-    CONF_SENSOR_PORT_TRAFFIC,
-    DEFAULT_SENSOR_PORT_TRAFFIC,
-    DOMAIN,
-    ATTRIBUTION,
-)
 from .sensor_types import (
-    MikrotikSensorEntityDescription,
     SENSOR_TYPES,
+    SENSOR_SERVICES,
     DEVICE_ATTRIBUTES_IFACE_ETHER,
     DEVICE_ATTRIBUTES_IFACE_SFP,
 )
@@ -33,172 +20,28 @@ _LOGGER = logging.getLogger(__name__)
 #   async_setup_entry
 # ---------------------------
 async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up device tracker for Mikrotik Router component."""
+    """Set up entry for component"""
-    inst = config_entry.data[CONF_NAME]
+    dispatcher = {
-    mikrotik_controller = hass.data[DOMAIN][config_entry.entry_id]
+        "MikrotikSensor": MikrotikSensor,
-    sensors = {}
+        "MikrotikInterfaceTrafficSensor": MikrotikInterfaceTrafficSensor,
-
+        "MikrotikClientTrafficSensor": MikrotikClientTrafficSensor,
-    @callback
+    }
-    def update_controller():
+    await model_async_setup_entry(
-        """Update the values of the controller."""
+        hass,
-        update_items(
+        config_entry,
-            inst, config_entry, mikrotik_controller, async_add_entities, sensors
+        async_add_entities,
-        )
+        SENSOR_SERVICES,
-
+        SENSOR_TYPES,
-    mikrotik_controller.listeners.append(
+        dispatcher,
-        async_dispatcher_connect(
-            hass, mikrotik_controller.signal_update, update_controller
-        )
     )
 
-    update_controller()
-
 
 # ---------------------------
-#   update_items
+#   MikrotikSensor
 # ---------------------------
-@callback
+class MikrotikSensor(MikrotikEntity, SensorEntity):
-def update_items(inst, config_entry, mikrotik_controller, async_add_entities, sensors):
-    """Update sensor state from the controller."""
-    new_sensors = []
-
-    for sensor, sid_func in zip(
-        # Sensor type name
-        [
-            "environment",
-            "traffic_rx",
-            "traffic_tx",
-            "client_traffic_rx",
-            "client_traffic_tx",
-            "client_traffic_lan_rx",
-            "client_traffic_lan_tx",
-            "client_traffic_wan_rx",
-            "client_traffic_wan_tx",
-        ],
-        # Entity function
-        [
-            MikrotikControllerSensor,
-            MikrotikInterfaceTrafficSensor,
-            MikrotikInterfaceTrafficSensor,
-            MikrotikClientTrafficSensor,
-            MikrotikClientTrafficSensor,
-            MikrotikClientTrafficSensor,
-            MikrotikClientTrafficSensor,
-            MikrotikClientTrafficSensor,
-            MikrotikClientTrafficSensor,
-        ],
-    ):
-        if sensor.startswith("traffic_") and not config_entry.options.get(
-            CONF_SENSOR_PORT_TRAFFIC, DEFAULT_SENSOR_PORT_TRAFFIC
-        ):
-            continue
-
-        uid_sensor = SENSOR_TYPES[sensor]
-        for uid in mikrotik_controller.data[uid_sensor.data_path]:
-            uid_data = mikrotik_controller.data[uid_sensor.data_path]
-            if (
-                uid_sensor.data_path == "interface"
-                and uid_data[uid]["type"] == "bridge"
-            ):
-                continue
-
-            if (
-                uid_sensor.data_path == "client_traffic"
-                and uid_sensor.data_attribute not in uid_data[uid].keys()
-            ):
-                continue
-
-            item_id = f"{inst}-{sensor}-{uid_data[uid][uid_sensor.data_reference]}"
-            _LOGGER.debug("Updating sensor %s", item_id)
-            if item_id in sensors:
-                if sensors[item_id].enabled:
-                    sensors[item_id].async_schedule_update_ha_state()
-                continue
-
-            sensors[item_id] = sid_func(
-                inst=inst,
-                uid=uid,
-                mikrotik_controller=mikrotik_controller,
-                entity_description=uid_sensor,
-            )
-            new_sensors.append(sensors[item_id])
-
-    for sensor in SENSOR_TYPES:
-        if sensor.startswith("system_"):
-            uid_sensor = SENSOR_TYPES[sensor]
-            if (
-                uid_sensor.data_attribute
-                not in mikrotik_controller.data[uid_sensor.data_path]
-                or mikrotik_controller.data[uid_sensor.data_path][
-                    uid_sensor.data_attribute
-                ]
-                == "unknown"
-            ):
-                continue
-            item_id = f"{inst}-{sensor}"
-            _LOGGER.debug("Updating sensor %s", item_id)
-            if item_id in sensors:
-                if sensors[item_id].enabled:
-                    sensors[item_id].async_schedule_update_ha_state()
-                continue
-
-            sensors[item_id] = MikrotikControllerSensor(
-                inst=inst,
-                uid="",
-                mikrotik_controller=mikrotik_controller,
-                entity_description=uid_sensor,
-            )
-            new_sensors.append(sensors[item_id])
-
-    if new_sensors:
-        async_add_entities(new_sensors, True)
-
-
-# ---------------------------
-#   MikrotikControllerSensor
-# ---------------------------
-class MikrotikControllerSensor(SensorEntity):
     """Define an Mikrotik Controller sensor."""
 
-    def __init__(
-        self,
-        inst,
-        uid: "",
-        mikrotik_controller,
-        entity_description: MikrotikSensorEntityDescription,
-    ):
-        """Initialize."""
-        self.entity_description = entity_description
-        self._inst = inst
-        self._ctrl = mikrotik_controller
-        self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: ATTRIBUTION}
-        self._uid = uid
-        if self._uid:
-            self._data = mikrotik_controller.data[self.entity_description.data_path][
-                self._uid
-            ]
-        else:
-            self._data = mikrotik_controller.data[self.entity_description.data_path]
-
-    @property
-    def name(self) -> str:
-        """Return the name."""
-        if self._uid:
-            if self.entity_description.name:
-                return f"{self._inst} {self._data[self.entity_description.data_name]} {self.entity_description.name}"
-
-            return f"{self._inst} {self._data[self.entity_description.data_name]}"
-        else:
-            return f"{self._inst} {self.entity_description.name}"
-
-    @property
-    def unique_id(self) -> str:
-        """Return a unique id for this entity."""
-        if self._uid:
-            return f"{self._inst.lower()}-{self.entity_description.key}-{self._data[self.entity_description.data_reference].lower()}"
-        else:
-            return f"{self._inst.lower()}-{self.entity_description.key}"
-
     @property
     def state(self) -> Optional[str]:
         """Return the state."""
@@ -221,87 +64,11 @@ class MikrotikControllerSensor(SensorEntity):
 
         return None
 
-    @property
-    def available(self) -> bool:
-        """Return if controller is available."""
-        return self._ctrl.connected()
-
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return a description for device registry."""
-        dev_connection = DOMAIN
-        dev_connection_value = self.entity_description.data_reference
-        dev_group = self.entity_description.ha_group
-        if self.entity_description.ha_group == "System":
-            dev_group = self._ctrl.data["resource"]["board-name"]
-            dev_connection_value = self._ctrl.data["routerboard"]["serial-number"]
-
-        if self.entity_description.ha_group.startswith("data__"):
-            dev_group = self.entity_description.ha_group[6:]
-            if dev_group in self._data:
-                dev_group = self._data[dev_group]
-                dev_connection_value = dev_group
-
-        if self.entity_description.ha_connection:
-            dev_connection = self.entity_description.ha_connection
-
-        if self.entity_description.ha_connection_value:
-            dev_connection_value = self.entity_description.ha_connection_value
-            if dev_connection_value.startswith("data__"):
-                dev_connection_value = dev_connection_value[6:]
-                dev_connection_value = self._data[dev_connection_value]
-
-        info = DeviceInfo(
-            connections={(dev_connection, f"{dev_connection_value}")},
-            identifiers={(dev_connection, f"{dev_connection_value}")},
-            default_name=f"{self._inst} {dev_group}",
-            default_model=f"{self._ctrl.data['resource']['board-name']}",
-            deafult_manufacturer=f"{self._ctrl.data['resource']['platform']}",
-            sw_version=f"{self._ctrl.data['resource']['version']}",
-            configuration_url=f"http://{self._ctrl.config_entry.data[CONF_HOST]}",
-            via_device=(DOMAIN, f"{self._ctrl.data['routerboard']['serial-number']}"),
-        )
-
-        if "mac-address" in self.entity_description.data_reference:
-            dev_group = self._data[self.entity_description.data_name]
-            dev_manufacturer = ""
-            if dev_connection_value in self._ctrl.data["host"]:
-                dev_group = self._ctrl.data["host"][dev_connection_value]["host-name"]
-                dev_manufacturer = self._ctrl.data["host"][dev_connection_value][
-                    "manufacturer"
-                ]
-
-            info = DeviceInfo(
-                connections={(dev_connection, f"{dev_connection_value}")},
-                default_name=f"{dev_group}",
-                default_manufacturer=f"{dev_manufacturer}",
-                via_device=(
-                    DOMAIN,
-                    f"{self._ctrl.data['routerboard']['serial-number']}",
-                ),
-            )
-
-        return info
-
-    @property
-    def extra_state_attributes(self) -> Mapping[str, Any]:
-        """Return the state attributes."""
-        attributes = super().extra_state_attributes
-        for variable in self.entity_description.data_attributes_list:
-            if variable in self._data:
-                attributes[format_attribute(variable)] = self._data[variable]
-
-        return attributes
-
-    async def async_added_to_hass(self):
-        """Run when entity about to be added to hass."""
-        _LOGGER.debug("New sensor %s (%s)", self._inst, self.unique_id)
-
 
 # ---------------------------
 #   MikrotikInterfaceTrafficSensor
 # ---------------------------
-class MikrotikInterfaceTrafficSensor(MikrotikControllerSensor):
+class MikrotikInterfaceTrafficSensor(MikrotikSensor):
     """Define an Mikrotik MikrotikInterfaceTrafficSensor sensor."""
 
     @property
@@ -325,7 +92,7 @@ class MikrotikInterfaceTrafficSensor(MikrotikControllerSensor):
 # ---------------------------
 #   MikrotikClientTrafficSensor
 # ---------------------------
-class MikrotikClientTrafficSensor(MikrotikControllerSensor):
+class MikrotikClientTrafficSensor(MikrotikSensor):
     """Define an Mikrotik MikrotikClientTrafficSensor sensor."""
 
     @property

custom_components/mikrotik_router/sensor_types.py

@@ -76,9 +76,11 @@ class MikrotikSensorEntityDescription(SensorEntityDescription):
     data_path: str = ""
     data_attribute: str = ""
     data_name: str = ""
+    data_name_comment: bool = False
     data_uid: str = ""
     data_reference: str = ""
     data_attributes_list: List = field(default_factory=lambda: [])
+    func: str = "MikrotikSensor"
 
 
 SENSOR_TYPES = {
@@ -294,6 +296,7 @@ SENSOR_TYPES = {
         data_uid="",
         data_reference="default-name",
         data_attributes_list=DEVICE_ATTRIBUTES_IFACE,
+        func="MikrotikInterfaceTrafficSensor",
     ),
     "traffic_rx": MikrotikSensorEntityDescription(
         key="traffic_rx",
@@ -312,6 +315,7 @@ SENSOR_TYPES = {
         data_uid="",
         data_reference="default-name",
         data_attributes_list=DEVICE_ATTRIBUTES_IFACE,
+        func="MikrotikInterfaceTrafficSensor",
     ),
     "client_traffic_lan_tx": MikrotikSensorEntityDescription(
         key="client_traffic_lan_tx",
@@ -330,6 +334,7 @@ SENSOR_TYPES = {
         data_uid="",
         data_reference="mac-address",
         data_attributes_list=DEVICE_ATTRIBUTES_CLIENT_TRAFFIC,
+        func="MikrotikClientTrafficSensor",
     ),
     "client_traffic_lan_rx": MikrotikSensorEntityDescription(
         key="client_traffic_lan_rx",
@@ -348,6 +353,7 @@ SENSOR_TYPES = {
         data_uid="",
         data_reference="mac-address",
         data_attributes_list=DEVICE_ATTRIBUTES_CLIENT_TRAFFIC,
+        func="MikrotikClientTrafficSensor",
     ),
     "client_traffic_wan_tx": MikrotikSensorEntityDescription(
         key="client_traffic_wan_tx",
@@ -366,6 +372,7 @@ SENSOR_TYPES = {
         data_uid="",
         data_reference="mac-address",
         data_attributes_list=DEVICE_ATTRIBUTES_CLIENT_TRAFFIC,
+        func="MikrotikClientTrafficSensor",
     ),
     "client_traffic_wan_rx": MikrotikSensorEntityDescription(
         key="client_traffic_wan_rx",
@@ -384,6 +391,7 @@ SENSOR_TYPES = {
         data_uid="",
         data_reference="mac-address",
         data_attributes_list=DEVICE_ATTRIBUTES_CLIENT_TRAFFIC,
+        func="MikrotikClientTrafficSensor",
     ),
     "client_traffic_tx": MikrotikSensorEntityDescription(
         key="client_traffic_tx",
@@ -402,6 +410,7 @@ SENSOR_TYPES = {
         data_uid="",
         data_reference="mac-address",
         data_attributes_list=DEVICE_ATTRIBUTES_CLIENT_TRAFFIC,
+        func="MikrotikClientTrafficSensor",
     ),
     "client_traffic_rx": MikrotikSensorEntityDescription(
         key="client_traffic_rx",
@@ -420,6 +429,7 @@ SENSOR_TYPES = {
         data_uid="",
         data_reference="mac-address",
         data_attributes_list=DEVICE_ATTRIBUTES_CLIENT_TRAFFIC,
+        func="MikrotikClientTrafficSensor",
     ),
     "environment": MikrotikSensorEntityDescription(
         key="environment",
@@ -439,3 +449,5 @@ SENSOR_TYPES = {
         data_reference="name",
     ),
 }
+
+SENSOR_SERVICES = {}

custom_components/mikrotik_router/switch.py

@@ -4,16 +4,12 @@ import logging
 from typing import Any, Optional
 from collections.abc import Mapping
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.const import CONF_NAME, CONF_HOST, ATTR_ATTRIBUTION
-from homeassistant.core import callback
-from homeassistant.helpers.entity import DeviceInfo
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.restore_state import RestoreEntity
 from .helper import format_attribute
-from .const import DOMAIN, ATTRIBUTION
+from .model import model_async_setup_entry, MikrotikEntity
 from .switch_types import (
-    MikrotikSwitchEntityDescription,
+    SENSOR_TYPES,
-    SWITCH_TYPES,
+    SENSOR_SERVICES,
     DEVICE_ATTRIBUTES_IFACE_ETHER,
     DEVICE_ATTRIBUTES_IFACE_SFP,
 )
@@ -25,142 +21,45 @@ _LOGGER = logging.getLogger(__name__)
 #   async_setup_entry
 # ---------------------------
 async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up switches for Mikrotik Router component."""
+    """Set up entry for component"""
-    inst = config_entry.data[CONF_NAME]
+    dispatcher = {
-    mikrotik_controller = hass.data[DOMAIN][config_entry.entry_id]
+        "MikrotikSwitch": MikrotikSwitch,
-    switches = {}
+        "MikrotikPortSwitch": MikrotikPortSwitch,
-
+        "MikrotikNATSwitch": MikrotikNATSwitch,
-    @callback
+        "MikrotikMangleSwitch": MikrotikMangleSwitch,
-    def update_controller():
+        "MikrotikFilterSwitch": MikrotikFilterSwitch,
-        """Update the values of the controller."""
+        "MikrotikQueueSwitch": MikrotikQueueSwitch,
-        update_items(inst, mikrotik_controller, async_add_entities, switches)
+        "MikrotikKidcontrolPauseSwitch": MikrotikKidcontrolPauseSwitch,
-
+    }
-    mikrotik_controller.listeners.append(
+    await model_async_setup_entry(
-        async_dispatcher_connect(
+        hass,
-            hass, mikrotik_controller.signal_update, update_controller
+        config_entry,
-        )
+        async_add_entities,
+        SENSOR_SERVICES,
+        SENSOR_TYPES,
+        dispatcher,
     )
 
-    update_controller()
-
 
 # ---------------------------
-#   update_items
+#   MikrotikSwitch
 # ---------------------------
-@callback
+class MikrotikSwitch(MikrotikEntity, SwitchEntity, RestoreEntity):
-def update_items(inst, mikrotik_controller, async_add_entities, switches):
-    """Update device switch state from the controller."""
-    new_switches = []
-
-    # Add switches
-    for switch, sid_func in zip(
-        # Switch type name
-        [
-            "interface",
-            "nat",
-            "mangle",
-            "filter",
-            "ppp_secret",
-            "queue",
-            "kidcontrol_enable",
-            "kidcontrol_pause",
-        ],
-        # Entity function
-        [
-            MikrotikControllerPortSwitch,
-            MikrotikControllerNATSwitch,
-            MikrotikControllerMangleSwitch,
-            MikrotikControllerFilterSwitch,
-            MikrotikControllerSwitch,
-            MikrotikControllerQueueSwitch,
-            MikrotikControllerSwitch,
-            MikrotikControllerKidcontrolPauseSwitch,
-        ],
-    ):
-        uid_switch = SWITCH_TYPES[switch]
-        for uid in mikrotik_controller.data[SWITCH_TYPES[switch].data_path]:
-            uid_data = mikrotik_controller.data[SWITCH_TYPES[switch].data_path]
-            item_id = f"{inst}-{switch}-{uid_data[uid][uid_switch.data_reference]}"
-            _LOGGER.debug("Updating sensor %s", item_id)
-            if item_id in switches:
-                if switches[item_id].enabled:
-                    switches[item_id].async_schedule_update_ha_state()
-                continue
-
-            switches[item_id] = sid_func(
-                inst=inst,
-                uid=uid,
-                mikrotik_controller=mikrotik_controller,
-                entity_description=uid_switch,
-            )
-            new_switches.append(switches[item_id])
-
-    if new_switches:
-        async_add_entities(new_switches)
-
-
-# ---------------------------
-#   MikrotikControllerSwitch
-# ---------------------------
-class MikrotikControllerSwitch(SwitchEntity, RestoreEntity):
     """Representation of a switch."""
 
-    def __init__(
-        self,
-        inst,
-        uid,
-        mikrotik_controller,
-        entity_description: MikrotikSwitchEntityDescription,
-    ):
-        self.entity_description = entity_description
-        self._inst = inst
-        self._ctrl = mikrotik_controller
-        self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: ATTRIBUTION}
-        self._data = mikrotik_controller.data[self.entity_description.data_path][uid]
-
-    @property
-    def available(self) -> bool:
-        """Return if controller is available."""
-        return self._ctrl.connected()
-
-    @property
-    def name(self) -> str:
-        """Return the name."""
-        if self.entity_description.data_name_comment and self._data["comment"]:
-            return (
-                f"{self._inst} {self.entity_description.name} {self._data['comment']}"
-            )
-
-        return f"{self._inst} {self.entity_description.name} {self._data[self.entity_description.data_name]}"
-
-    @property
-    def unique_id(self) -> str:
-        """Return a unique id for this entity."""
-        return f"{self._inst.lower()}-{self.entity_description.key}-{self._data[self.entity_description.data_reference].lower()}"
-
     @property
     def is_on(self) -> bool:
         """Return true if device is on."""
-        return self._data[self.entity_description.data_is_on]
+        return self._data[self.entity_description.data_attribute]
 
     @property
     def icon(self) -> str:
         """Return the icon."""
-        if self._data[self.entity_description.data_is_on]:
+        if self._data[self.entity_description.data_attribute]:
             return self.entity_description.icon_enabled
         else:
             return self.entity_description.icon_disabled
 
-    @property
-    def extra_state_attributes(self) -> Mapping[str, Any]:
-        """Return the state attributes."""
-        attributes = super().extra_state_attributes
-        for variable in self.entity_description.data_attributes_list:
-            if variable in self._data:
-                attributes[format_attribute(variable)] = self._data[variable]
-
-        return attributes
-
     def turn_on(self, **kwargs: Any) -> None:
         """Required abstract method."""
         pass
@@ -187,49 +86,11 @@ class MikrotikControllerSwitch(SwitchEntity, RestoreEntity):
         self._ctrl.set_value(path, param, value, mod_param, True)
         await self._ctrl.async_update()
 
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return a description for device registry."""
-        dev_connection = DOMAIN
-        dev_connection_value = self.entity_description.data_reference
-        dev_group = self.entity_description.ha_group
-        if self.entity_description.ha_group.startswith("data__"):
-            dev_group = self.entity_description.ha_group[6:]
-            if dev_group in self._data:
-                dev_group = self._data[dev_group]
-                dev_connection_value = dev_group
-
-        if self.entity_description.ha_connection:
-            dev_connection = self.entity_description.ha_connection
-
-        if self.entity_description.ha_connection_value:
-            dev_connection_value = self.entity_description.ha_connection_value
-            if dev_connection_value.startswith("data__"):
-                dev_connection_value = dev_connection_value[6:]
-                dev_connection_value = self._data[dev_connection_value]
-
-        info = DeviceInfo(
-            connections={(dev_connection, f"{dev_connection_value}")},
-            identifiers={(dev_connection, f"{dev_connection_value}")},
-            default_name=f"{self._inst} {dev_group}",
-            default_model=f"{self._ctrl.data['resource']['board-name']}",
-            default_manufacturer=f"{self._ctrl.data['resource']['platform']}",
-            sw_version=f"{self._ctrl.data['resource']['version']}",
-            configuration_url=f"http://{self._ctrl.config_entry.data[CONF_HOST]}",
-            via_device=(DOMAIN, f"{self._ctrl.data['routerboard']['serial-number']}"),
-        )
-
-        return info
-
-    async def async_added_to_hass(self):
-        """Run when entity about to be added to hass."""
-        _LOGGER.debug("New switch %s (%s)", self._inst, self.unique_id)
-
 
 # ---------------------------
-#   MikrotikControllerPortSwitch
+#   MikrotikPortSwitch
 # ---------------------------
-class MikrotikControllerPortSwitch(MikrotikControllerSwitch):
+class MikrotikPortSwitch(MikrotikSwitch):
     """Representation of a network port switch."""
 
     @property
@@ -302,9 +163,9 @@ class MikrotikControllerPortSwitch(MikrotikControllerSwitch):
 
 
 # ---------------------------
-#   MikrotikControllerNATSwitch
+#   MikrotikNATSwitch
 # ---------------------------
-class MikrotikControllerNATSwitch(MikrotikControllerSwitch):
+class MikrotikNATSwitch(MikrotikSwitch):
     """Representation of a NAT switch."""
 
     @property
@@ -351,9 +212,9 @@ class MikrotikControllerNATSwitch(MikrotikControllerSwitch):
 
 
 # ---------------------------
-#   MikrotikControllerMangleSwitch
+#   MikrotikMangleSwitch
 # ---------------------------
-class MikrotikControllerMangleSwitch(MikrotikControllerSwitch):
+class MikrotikMangleSwitch(MikrotikSwitch):
     """Representation of a Mangle switch."""
 
     async def async_turn_on(self) -> None:
@@ -394,9 +255,9 @@ class MikrotikControllerMangleSwitch(MikrotikControllerSwitch):
 
 
 # ---------------------------
-#   MikrotikControllerFilterSwitch
+#   MikrotikFilterSwitch
 # ---------------------------
-class MikrotikControllerFilterSwitch(MikrotikControllerSwitch):
+class MikrotikFilterSwitch(MikrotikSwitch):
     """Representation of a Filter switch."""
 
     async def async_turn_on(self) -> None:
@@ -435,9 +296,9 @@ class MikrotikControllerFilterSwitch(MikrotikControllerSwitch):
 
 
 # ---------------------------
-#   MikrotikControllerQueueSwitch
+#   MikrotikQueueSwitch
 # ---------------------------
-class MikrotikControllerQueueSwitch(MikrotikControllerSwitch):
+class MikrotikQueueSwitch(MikrotikSwitch):
     """Representation of a queue switch."""
 
     async def async_turn_on(self) -> None:
@@ -468,9 +329,9 @@ class MikrotikControllerQueueSwitch(MikrotikControllerSwitch):
 
 
 # ---------------------------
-#   MikrotikControllerKidcontrolPauseSwitch
+#   MikrotikKidcontrolPauseSwitch
 # ---------------------------
-class MikrotikControllerKidcontrolPauseSwitch(MikrotikControllerSwitch):
+class MikrotikKidcontrolPauseSwitch(MikrotikSwitch):
     """Representation of a queue switch."""
 
     async def async_turn_on(self) -> None:

custom_components/mikrotik_router/switch_types.py

@@ -2,7 +2,6 @@
 from dataclasses import dataclass, field
 from typing import List
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
-from homeassistant.helpers.entity import EntityCategory
 from homeassistant.components.switch import (
     SwitchDeviceClass,
     SwitchEntityDescription,
@@ -147,7 +146,7 @@ class MikrotikSwitchEntityDescription(SwitchEntityDescription):
     ha_connection: str = ""
     ha_connection_value: str = ""
     data_path: str = ""
-    data_is_on: str = "enabled"
+    data_attribute: str = "enabled"
     data_switch_path: str = ""
     data_switch_parameter: str = "disabled"
     data_name: str = ""
@@ -155,9 +154,10 @@ class MikrotikSwitchEntityDescription(SwitchEntityDescription):
     data_uid: str = ""
     data_reference: str = ""
     data_attributes_list: List = field(default_factory=lambda: [])
+    func: str = "MikrotikSwitch"
 
 
-SWITCH_TYPES = {
+SENSOR_TYPES = {
     "interface": MikrotikSwitchEntityDescription(
         key="interface",
         name="port",
@@ -173,6 +173,7 @@ SWITCH_TYPES = {
         data_uid="name",
         data_reference="default-name",
         data_attributes_list=DEVICE_ATTRIBUTES_IFACE,
+        func="MikrotikPortSwitch",
     ),
     "nat": MikrotikSwitchEntityDescription(
         key="nat",
@@ -190,6 +191,7 @@ SWITCH_TYPES = {
         data_uid="uniq-id",
         data_reference="uniq-id",
         data_attributes_list=DEVICE_ATTRIBUTES_NAT,
+        func="MikrotikNATSwitch",
     ),
     "mangle": MikrotikSwitchEntityDescription(
         key="mangle",
@@ -207,6 +209,7 @@ SWITCH_TYPES = {
         data_uid="uniq-id",
         data_reference="uniq-id",
         data_attributes_list=DEVICE_ATTRIBUTES_MANGLE,
+        func="MikrotikMangleSwitch",
     ),
     "filter": MikrotikSwitchEntityDescription(
         key="filter",
@@ -224,6 +227,7 @@ SWITCH_TYPES = {
         data_uid="uniq-id",
         data_reference="uniq-id",
         data_attributes_list=DEVICE_ATTRIBUTES_FILTER,
+        func="MikrotikFilterSwitch",
     ),
     "ppp_secret": MikrotikSwitchEntityDescription(
         key="ppp_secret",
@@ -256,6 +260,7 @@ SWITCH_TYPES = {
         data_uid="name",
         data_reference="name",
         data_attributes_list=DEVICE_ATTRIBUTES_QUEUE,
+        func="MikrotikQueueSwitch",
     ),
     "kidcontrol_enable": MikrotikSwitchEntityDescription(
         key="kidcontrol_enable",
@@ -283,11 +288,14 @@ SWITCH_TYPES = {
         ha_connection=DOMAIN,
         ha_connection_value="Kidcontrol",
         data_path="kid-control",
-        data_is_on="paused",
+        data_attribute="paused",
         data_switch_path="/ip/kid-control",
         data_name="name",
         data_uid="name",
         data_reference="name",
         data_attributes_list=DEVICE_ATTRIBUTES_KIDCONTROL,
+        func="MikrotikKidcontrolPauseSwitch",
     ),
 }
+
+SENSOR_SERVICES = {}