Apply persisted role to messages missing role

Author: jmattsson

custom_components/powersensor/PowersensorMessageDispatcher.py

@@ -3,6 +3,7 @@
 import logging
 
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers.dispatcher import async_dispatcher_send, async_dispatcher_connect
 
 from powersensor_local import PlugApi, VirtualHousehold
@@ -12,8 +13,9 @@
 
 _LOGGER = logging.getLogger(__name__)
 class PowersensorMessageDispatcher:
-    def __init__(self, hass: HomeAssistant, vhh: VirtualHousehold, debounce_timeout: float = 60):
+    def __init__(self, hass: HomeAssistant, entry: ConfigEntry, vhh: VirtualHousehold, debounce_timeout: float = 60):
         self._hass = hass
+        self._entry = entry
         self._vhh = vhh
         self.plugs = dict()
         self._known_plugs = set()
@@ -163,6 +165,8 @@ def cancel_any_pending_removal(self, mac, source):
     async def handle_message(self, event: str, message: dict):
         mac = message['mac']
         role = message.get('role', None)
+        if role is None:
+            role = self._entry.data.get('roles', {}).get(mac, None)
         self.cancel_any_pending_removal(mac, "new message received from plug")
 
         if mac not in self.plugs.keys():

custom_components/powersensor/__init__.py

@@ -51,7 +51,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     vhh = VirtualHousehold(entry.data['with_solar'])
 
     # Set up message dispatcher
-    dispatcher = PowersensorMessageDispatcher(hass, vhh)
+    dispatcher = PowersensorMessageDispatcher(hass, entry, vhh)
     for mac, network_info in entry.data['devices'].items():
         await dispatcher.enqueue_plug_for_adding(network_info)