Add support for simple queues. Transfer traffic_type and traffic_div calculation to separate method

custom_components/mikrotik_router/mikrotik_controller.py

@@ -55,6 +55,7 @@ class MikrotikControllerData:
             "nat": {},
             "fw-update": {},
             "script": {},
+            "queue": {},
         }
 
         self.listeners = []
@@ -69,6 +70,23 @@ class MikrotikControllerData:
             self.hass, self.force_fwupdate_check, timedelta(hours=1)
         )
 
+    def _get_traffic_type_and_div(self):
+        traffic_type = self.option_traffic_type
+        if traffic_type == "Kbps":
+            traffic_div = 0.001
+        elif traffic_type == "Mbps":
+            traffic_div = 0.000001
+        elif traffic_type == "B/s":
+            traffic_div = 0.125
+        elif traffic_type == "KB/s":
+            traffic_div = 0.000125
+        elif traffic_type == "MB/s":
+            traffic_div = 0.000000125
+        else:
+            traffic_type = "bps"
+            traffic_div = 1
+        return traffic_type, traffic_div
+
     # ---------------------------
     #   force_update
     # ---------------------------
@@ -168,6 +186,7 @@ class MikrotikControllerData:
         await self.hass.async_add_executor_job(self.get_nat)
         await self.hass.async_add_executor_job(self.get_system_resource)
         await self.hass.async_add_executor_job(self.get_script)
+        await self.hass.async_add_executor_job(self.get_queue)
 
         async_dispatcher_send(self.hass, self.signal_update)
         self.lock.release()
@@ -257,20 +276,7 @@ class MikrotikControllerData:
             ],
         )
 
-        traffic_type = self.option_traffic_type
+        traffic_type, traffic_div = self._get_traffic_type_and_div()
-        if traffic_type == "Kbps":
-            traffic_div = 0.001
-        elif traffic_type == "Mbps":
-            traffic_div = 0.000001
-        elif traffic_type == "B/s":
-            traffic_div = 0.125
-        elif traffic_type == "KB/s":
-            traffic_div = 0.000125
-        elif traffic_type == "MB/s":
-            traffic_div = 0.000000125
-        else:
-            traffic_type = "bps"
-            traffic_div = 1
 
         for uid in self.data["interface"]:
             self.data["interface"][uid]["rx-bits-per-second-attr"] = traffic_type
@@ -559,3 +565,60 @@ class MikrotikControllerData:
                 {"name": "run-count", "default": "unknown"},
             ],
         )
+
+    # ---------------------------
+    #   get_queue
+    # ---------------------------
+
+    def get_queue(self):
+        """Get Queue data from Mikrotik"""
+        self.data["queue"] = parse_api(
+            data=self.data["queue"],
+            source=self.api.path("/queue/simple"),
+            key="name",
+            vals=[
+                {"name": "name", "default": "unknown"},
+                {"name": "target", "default": "unknown"},
+                {"name": "max-limit", "default": "0/0"},
+                {"name": "limit-at", "default": "0/0"},
+                {"name": "burst-limit", "default": "0/0"},
+                {"name": "burst-threshold", "default": "0/0"},
+                {"name": "burst-time", "default": "0s/0s"},
+                {"name": "packet-marks", "default": "none"},
+                {"name": "parent", "default": "none"},
+                {"name": "comment"},
+                {
+                    "name": "enabled",
+                    "source": "disabled",
+                    "type": "bool",
+                    "reverse": True,
+                },
+            ]
+        )
+
+        traffic_type, traffic_div = self._get_traffic_type_and_div()
+
+        for uid in self.data["queue"]:
+            upload_max_limit_bps, download_max_limit_bps = [int(x) for x in
+                                                            self.data["queue"][uid]["max-limit"].split('/')]
+            self.data["queue"][uid]["upload-max-limit"] = f"{round(upload_max_limit_bps * traffic_div)} {traffic_type}"
+            self.data["queue"][uid]["download-max-limit"] = f"{round(download_max_limit_bps * traffic_div)} {traffic_type}"
+
+            upload_limit_at_bps, download_limit_at_bps = [int(x) for x in
+                                                          self.data["queue"][uid]["limit-at"].split('/')]
+            self.data["queue"][uid]["upload-limit-at"] = f"{round(upload_limit_at_bps * traffic_div)} {traffic_type}"
+            self.data["queue"][uid]["download-limit-at"] = f"{round(download_limit_at_bps * traffic_div)} {traffic_type}"
+
+            upload_burst_limit_bps, download_burst_limit_bps = [int(x) for x in
+                                                                self.data["queue"][uid]["burst-limit"].split('/')]
+            self.data["queue"][uid]["upload-burst-limit"] = f"{round(upload_burst_limit_bps * traffic_div)} {traffic_type}"
+            self.data["queue"][uid]["download-burst-limit"] = f"{round(download_burst_limit_bps * traffic_div)} {traffic_type}"
+
+            upload_burst_threshold_bps, download_burst_threshold_bps = [int(x) for x in
+                                                                        self.data["queue"][uid]["burst-threshold"].split('/')]
+            self.data["queue"][uid]["upload-burst-threshold"] = f"{round(upload_burst_threshold_bps * traffic_div)} {traffic_type}"
+            self.data["queue"][uid]["download-burst-threshold"] = f"{round(download_burst_threshold_bps * traffic_div)} {traffic_type}"
+
+            upload_burst_time, download_burst_time = self.data["queue"][uid]["burst-time"].split('/')
+            self.data["queue"][uid]["upload-burst-time"] = upload_burst_time
+            self.data["queue"][uid]["download-burst-time"] = download_burst_time

custom_components/mikrotik_router/switch.py

@@ -49,6 +49,23 @@ DEVICE_ATTRIBUTES_SCRIPT = [
     "run-count",
 ]
 
+DEVICE_ATTRIBUTES_QUEUE = [
+    "target",
+    "download-max-limit",
+    "upload-max-limit",
+    "upload-limit-at",
+    "download-limit-at",
+    "upload-burst-limit",
+    "download-burst-limit",
+    "upload-burst-threshold",
+    "download-burst-threshold",
+    "upload-burst-time",
+    "download-burst-time",
+    "packet-marks",
+    "parent",
+    "comment",
+]
+
 
 # ---------------------------
 #   format_attribute
@@ -95,11 +112,12 @@ def update_items(inst, mikrotik_controller, async_add_entities, switches):
 
     # Add switches
     for sid, sid_func in zip(
-        ["interface", "nat", "script"],
+        ["interface", "nat", "script", "queue"],
         [
             MikrotikControllerPortSwitch,
             MikrotikControllerNATSwitch,
             MikrotikControllerScriptSwitch,
+            MikrotikControllerQueueSwitch,
         ],
     ):
         for uid in mikrotik_controller.data[sid]:
@@ -421,3 +439,112 @@ class MikrotikControllerScriptSwitch(MikrotikControllerSwitch):
     def is_on(self):
         """Return true if device is on."""
         return False
+
+
+# ---------------------------
+#   MikrotikControllerNATSwitch
+# ---------------------------
+class MikrotikControllerQueueSwitch(MikrotikControllerSwitch):
+    """Representation of a queue switch."""
+
+    def __init__(self, inst, uid, mikrotik_controller):
+        """Set up queue switch."""
+        super().__init__(inst, uid, mikrotik_controller)
+
+        self._data = mikrotik_controller.data["queue"][self._uid]
+        self._attrs = {
+            ATTR_ATTRIBUTION: ATTRIBUTION,
+        }
+
+    async def async_added_to_hass(self):
+        """Queue switch entity created."""
+        _LOGGER.debug("New queue switch %s (%s)", self._inst, self._data["name"])
+
+    @property
+    def name(self) -> str:
+        """Return the name of the queue switch."""
+        return f"{self._inst} Queue {self._data['name']}"
+
+    @property
+    def unique_id(self) -> str:
+        """Return a unique identifier for this queue switch."""
+        return f"{self._inst.lower()}-queue_switch-{self._data['name']}"
+
+    @property
+    def icon(self):
+        """Return the icon."""
+        if not self._data["enabled"]:
+            icon = "mdi:leaf-off"
+        else:
+            icon = "mdi:leaf"
+
+        return icon
+
+    @property
+    def device_info(self):
+        """Return a queue switch description for device registry."""
+        info = {
+            "identifiers": {
+                (
+                    DOMAIN,
+                    "serial-number",
+                    self._ctrl.data["routerboard"]["serial-number"],
+                    "switch",
+                    "Queue",
+                )
+            },
+            "manufacturer": self._ctrl.data["resource"]["platform"],
+            "model": self._ctrl.data["resource"]["board-name"],
+            "name": "Queue",
+        }
+        return info
+
+    @property
+    def device_state_attributes(self):
+        """Return the queue switch state attributes."""
+        attributes = self._attrs
+
+        for variable in DEVICE_ATTRIBUTES_QUEUE:
+            if variable in self._data:
+                attributes[format_attribute(variable)] = self._data[variable]
+
+        return attributes
+
+    async def async_turn_on(self):
+        """Turn on the queue switch."""
+        path = "/queue/simple"
+        param = ".id"
+        value = None
+        for uid in self._ctrl.data["queue"]:
+            if (
+                self._ctrl.data["queue"][uid]["name"]
+                == f"{self._data['name']}"
+            ):
+                value = self._ctrl.data["queue"][uid][".id"]
+
+        mod_param = "disabled"
+        mod_value = False
+        self._ctrl.set_value(path, param, value, mod_param, mod_value)
+        await self._ctrl.force_update()
+
+    async def async_turn_off(self):
+        """Turn on the queue switch."""
+        path = "/queue/simple"
+        param = ".id"
+        value = None
+        for uid in self._ctrl.data["queue"]:
+            if (
+                self._ctrl.data["queue"][uid]["name"]
+                == f"{self._data['name']}"
+            ):
+                value = self._ctrl.data["queue"][uid][".id"]
+
+        mod_param = "disabled"
+        mod_value = True
+        self._ctrl.set_value(path, param, value, mod_param, mod_value)
+        await self._ctrl.async_update()
+
+    @property
+    def is_on(self):
+        """Return true if the queue is on."""
+        return self._data["enabled"]