added support for remapping keys for existing data structure

custom_components/mikrotik_router/helper.py

@@ -37,22 +37,24 @@ def from_entry_bool(entry, param, default=False, reverse=False):
 # ---------------------------
 #   from_list
 # ---------------------------
-async def from_list(data=None, source=None, key=None, vals=[], ensure_vals=[]):
+async def from_list(data=None, source=None, key=None, key_search=None, vals=[], ensure_vals=[]):
     if not source:
         return data
 
+    keymap = generate_keymap(data, key_search)
+
     for entry in source:
-        if key not in entry:
+        uid = await get_uid(entry, key)
+        if keymap and key_search in entry and entry[key_search] in keymap:
+            uid = keymap[entry[key_search]]
+
+        if not uid:
             continue
 
-        if not entry[key]:
-            continue
-
-        _LOGGER.debug("Processing entry {}, entry {}".format(source, entry))
-        uid = entry[key]
         if uid not in data:
             data[uid] = {}
 
+        _LOGGER.debug("Processing entry {}, entry {}".format(source, entry))
         for val in vals:
             _name = val['name']
             _type = val['type'] if 'type' in val else 'str'
@@ -75,3 +77,32 @@ async def from_list(data=None, source=None, key=None, vals=[], ensure_vals=[]):
                 data[uid][val['name']] = _default
 
     return data
+
+
+# ---------------------------
+#   get_uid
+# ---------------------------
+async def get_uid(entry, key):
+    if key not in entry:
+        return False
+
+    if not entry[key]:
+        return False
+
+    return entry[key]
+
+
+# ---------------------------
+#   generate_keymap
+# ---------------------------
+async def generate_keymap(data, key_search):
+    if not key_search:
+        return None
+
+    for uid in data:
+        if key_search not in uid:
+            continue
+
+        keymap[data[uid]['name']] = data[uid]['default-name']
+
+    return keymap

custom_components/mikrotik_router/mikrotik_controller.py

@@ -33,7 +33,6 @@ class MikrotikControllerData():
         self.data = {'routerboard': {},
                      'resource': {},
                      'interface': {},
-                     'interface_map': {},
                      'arp': {},
                      'nat': {},
                      'fw-update': {},
@@ -191,15 +190,6 @@ class MikrotikControllerData():
             ]
         )
 
-        interface_list = ""
-        for uid in self.data['interface']:
-            self.data['interface_map'][self.data['interface'][uid]['name']] = self.data['interface'][uid]['default-name']
-
-            if interface_list:
-                interface_list += ","
-
-            interface_list += self.data['interface'][uid]['name']
-
         await self.get_interface_traffic(interface_list)
         return
 
@@ -207,17 +197,23 @@ class MikrotikControllerData():
     #   get_interface_traffic
     # ---------------------------
     async def get_interface_traffic(self, interface_list):
-        data = await self.hass.async_add_executor_job(self.api.get_traffic, interface_list)
+        """Get traffic for all interfaces from Mikrotik"""
-        for entry in data:
+        interface_list = ""
-            iface_name = from_entry(entry, 'name')
+        for uid in self.data['interface']:
-            if iface_name not in self.data['interface_map']:
+            if interface_list:
-                continue
+                interface_list += ","
 
-            _LOGGER.debug("Processing entry {}, entry {}".format("/interface/monitor-traffic", entry))
+            interface_list += self.data['interface'][uid]['name']
-            uid = self.data['interface_map'][iface_name]
-            self.data['interface'][uid]['rx-bits-per-second'] = from_entry(entry, 'rx-bits-per-second', default=0)
-            self.data['interface'][uid]['tx-bits-per-second'] = from_entry(entry, 'tx-bits-per-second', default=0)
 
+        self.data['interface'] = await from_list(
+            data=self.data['interface'],
+            source=await self.hass.async_add_executor_job(self.api.get_traffic, interface_list),
+            key_search='name',
+            vals=[
+                {'name': 'rx-bits-per-second', 'default': 0},
+                {'name': 'tx-bits-per-second', 'default': 0},
+            ]
+        )
         return
 
     # ---------------------------