trying to rediscover ips missed by HA's zeroconf discovery

Author: bookman-dius

custom_components/powersensor/AsyncSet.py

@@ -0,0 +1,50 @@
+import asyncio
+
+class AsyncSet:
+    """Thread/async-safe set."""
+    
+    def __init__(self):
+        self._items = set()
+        self._lock = asyncio.Lock()
+    
+    async def add(self, item):
+        """Add item to set."""
+        async with self._lock:
+            self._items.add(item)
+    
+    async def discard(self, item):
+        """Remove item if present."""
+        async with self._lock:
+            self._items.discard(item)
+    
+    async def remove(self, item):
+        """Remove item, raise KeyError if not present."""
+        async with self._lock:
+            self._items.remove(item)
+    
+    async def pop(self):
+        """Remove and return arbitrary item."""
+        async with self._lock:
+            return self._items.pop()
+    
+    async def clear(self):
+        """Remove all items."""
+        async with self._lock:
+            self._items.clear()
+    
+    async def copy(self):
+        """Return a copy of the set."""
+        async with self._lock:
+            return self._items.copy()
+    
+    def __contains__(self, item):
+        """Check if item in set."""
+        return item in self._items
+    
+    def __len__(self):
+        """Get size of set."""
+        return len(self._items)
+
+    def __bool__(self):
+        """Check if set is empty in pythonic way"""
+        return bool(self._items)

custom_components/powersensor/PowersensorEntity.py

@@ -0,0 +1,133 @@
+from datetime import timedelta
+from typing import Callable
+
+from homeassistant.components.sensor import SensorEntity
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.event import async_track_point_in_utc_time
+from homeassistant.util.dt import utcnow
+from .PlugMeasurements import PlugMeasurements
+from .SensorMeasurements import SensorMeasurements
+from .const import DOMAIN, POWER_SENSOR_UPDATE_SIGNAL
+
+import logging
+_LOGGER = logging.getLogger(__name__)
+
+class PowersensorEntity(SensorEntity):
+    """Powersensor Plug Class--designed to handle all measurements of the plug--perhaps less expressive"""
+    def __init__(self, hass: HomeAssistant, mac : str,
+                 get_config_by_measurement_type: Callable[[SensorMeasurements|PlugMeasurements], dict],
+                 measurement_type: SensorMeasurements|PlugMeasurements, timeout_seconds: int = 60):
+        """Initialize the sensor."""
+        self.role = None
+        self._has_recently_received_update_message = False
+        self._attr_native_value = 0.0
+        self._hass = hass
+        self._mac = mac
+        self._model = f"PowersensorDevice"
+        self._device_name = f'PowersensorDevice MAC address: ({self._mac})'
+        self._measurement_name= None
+        self._remove_unavailability_tracker = None
+        self._timeout = timedelta(seconds=timeout_seconds)  # Adjust as needed
+
+        self.measurement_type = measurement_type
+        config = get_config_by_measurement_type(measurement_type)
+        self._attr_unique_id = f"powersensor_{mac}_{measurement_type}"
+        self._attr_device_class = config["device_class"]
+        self._attr_native_unit_of_measurement = config["unit"]
+        self._attr_device_info = self.device_info
+        self._attr_suggested_display_precision = config["precision"]
+        self._signal = f"{POWER_SENSOR_UPDATE_SIGNAL}_{self._mac}_{config['event']}"
+        if 'state_class' in config.keys():
+            self._attr_state_class = config['state_class']
+        self._message_key = config.get('message_key', None)
+        self._message_callback = config.get('callback', None)
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        raise NotImplementedError
+
+    @property
+    def available(self) -> bool:
+        """Does data exist for this sensor type"""
+        return self._has_recently_received_update_message
+
+    def _schedule_unavailable(self):
+        """Schedule entity to become unavailable."""
+        if self._remove_unavailability_tracker:
+            self._remove_unavailability_tracker()
+
+        self._remove_unavailability_tracker = async_track_point_in_utc_time(
+            self.hass,
+            self._async_make_unavailable,
+            utcnow() + self._timeout
+        )
+
+    async def _async_make_unavailable(self, _now):
+        """Mark entity as unavailable."""
+        self._has_recently_received_update_message = False
+        self.async_write_ha_state()
+
+    async def async_added_to_hass(self) -> None:
+        """Subscribe to messages when added to home assistant"""
+        self._has_recently_received_update_message = False
+        self.async_on_remove(async_dispatcher_connect(
+            self._hass,
+            self._signal,
+            self._handle_update
+        ))
+
+    async def async_will_remove_from_hass(self):
+        """Clean up."""
+        if self._remove_unavailability_tracker:
+            self._remove_unavailability_tracker()
+
+    def _rename_based_on_role(self):
+        return False
+
+    @callback
+    def _handle_update(self, event, message):
+        """handle pushed data."""
+        name_updated = False
+        if event =="remove_sensor":
+            self.async_remove()
+        elif event =="mark_unavailable":
+            self._has_recently_received_update_message = False
+        else:
+            self._has_recently_received_update_message = True
+            if not self.role:
+                if 'role' in message.keys():
+                    self.role = message['role']
+                    name_updated = self._rename_based_on_role()
+
+
+            if self._message_key in message.keys():
+                if self._message_callback:
+                    self._attr_native_value = self._message_callback( message[self._message_key])
+                else:
+                    self._attr_native_value = message[self._message_key]
+            self._schedule_unavailable()
+
+            if name_updated:
+                device_registry = dr.async_get(self.hass)
+                device = device_registry.async_get_device(
+                    identifiers={(DOMAIN, self._mac)}
+                )
+
+                if device and device.name != self._device_name:
+                    # Update the device name
+                    device_registry.async_update_device(
+                        device.id,
+                        name=self._device_name
+                    )
+
+                entity_registry = er.async_get(self.hass)
+                entity_registry.async_update_entity(
+                    self.entity_id,
+                    name = self._attr_name
+                )
+            self.async_write_ha_state()
+

custom_components/powersensor/PowersensorMessageDispatacher.py

@@ -1,18 +1,21 @@
+import asyncio
 import logging
 
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_send, async_dispatcher_connect
 
 from powersensor_local import PlugApi, VirtualHousehold
 
-from custom_components.powersensor.const import POWER_SENSOR_UPDATE_SIGNAL, DOMAIN
+from custom_components.powersensor.AsyncSet import AsyncSet
+from custom_components.powersensor.const import POWER_SENSOR_UPDATE_SIGNAL, DOMAIN, DEFAULT_PORT
 
 _LOGGER = logging.getLogger(__name__)
 class PowersensorMessageDispatcher:
     def __init__(self, hass: HomeAssistant, vhh: VirtualHousehold):
         self._hass = hass
         self._vhh = vhh
         self.plugs = dict()
+        self._known_plugs = set()
         self.sensors = dict()
         self.on_start_sensor_queue = dict()
         self._unsubscribe_from_signals = [
@@ -27,15 +30,75 @@ def __init__(self, hass: HomeAssistant, vhh: VirtualHousehold):
                                      self._plug_updated),
             async_dispatcher_connect(self._hass,
                                      f"{DOMAIN}_zeroconf_remove_plug",
-                                     self._plug_remove)
+                                     self._plug_remove),
+            async_dispatcher_connect(self._hass,
+                                     f"{DOMAIN}_plug_added_to_homeassistant",
+                                     self._acknowledge_plug_added_to_homeassistant),
         ]
 
-
-    def add_api(self, mac, network_info):
-
-        _LOGGER.info(f"Adding API for mac={network_info['mac']}, ip={network_info['host']}, port={network_info['port']}")
-        api = PlugApi(mac=network_info['mac'], ip=network_info['host'], port=network_info['port'])
-        self.plugs[mac] = api
+        self._monitor_add_plug_queue = None
+        self._stop_task = False
+        self._plug_added_queue = AsyncSet()
+        self._safe_to_process_plug_queue = False
+
+    async def enqueue_plug_for_adding(self, network_info: dict):
+        _LOGGER.debug(f"Adding to plug processing queue: {network_info}")
+        await self._plug_added_queue.add((network_info['mac'], network_info['host'], network_info['port']))
+
+    async def process_plug_queue(self):
+        """Start the background task if not already running."""
+        self._safe_to_process_plug_queue = True
+        if self._monitor_add_plug_queue is None or self._monitor_add_plug_queue.done():
+            self._stop_task = False
+            self._monitor_add_plug_queue = self._hass.async_create_background_task(self._monitor_plug_queue(), name="plug_queue_monitor")
+            _LOGGER.debug("Background task started")
+
+    def _plug_has_been_seen(self, mac_address)->bool:
+        return mac_address in self.plugs or mac_address in self._known_plugs
+
+    async def _monitor_plug_queue(self):
+        """The actual background task loop."""
+        try:
+            while not self._stop_task and self._plug_added_queue:
+                queue_snapshot = await self._plug_added_queue.copy()
+                for mac_address, host, port in queue_snapshot:
+                    if not self._plug_has_been_seen(mac_address):
+                        async_dispatcher_send(self._hass, f"{DOMAIN}_create_plug",
+                                              mac_address, host, port)
+                    else:
+                        _LOGGER.debug(f"Plug: {mac_address} has already been created as an entity in Home Assistant."
+                                      f" Skipping and flushing from queue.")
+                        await self._plug_added_queue.remove((mac_address, host, port))
+
+
+                await asyncio.sleep(5)
+            _LOGGER.debug("Plug queue has been processed!")
+
+        except asyncio.CancelledError:
+            _LOGGER.debug("Plug queue processing cancelled")
+            raise
+        except Exception as e:
+            _LOGGER.error(f"Error in Plug queue processing task: {e}")
+        finally:
+            self._monitor_add_plug_queue = None
+
+    async def stop_processing_plug_queue(self):
+        """Stop the background task."""
+        self._stop_task = True
+        if self._monitor_add_plug_queue and not self._monitor_add_plug_queue.done():
+            self._monitor_add_plug_queue.cancel()
+            try:
+                await self._monitor_add_plug_queue
+            except asyncio.CancelledError:
+                pass
+            _LOGGER.debug("Background task stopped")
+            self._monitor_add_plug_queue = None
+
+    def _create_api(self, mac_address, ip, port):
+        _LOGGER.info(f"Adding API for mac={mac_address}, ip={ip}, port={port}")
+        api = PlugApi(mac=mac_address, ip=ip, port=port)
+        self.plugs[mac_address] = api
+        self._known_plugs.add(mac_address)
         known_evs = [
             'exception',
             'average_flow',
@@ -53,6 +116,10 @@ def add_api(self, mac, network_info):
             api.subscribe(ev, lambda event, message: self.handle_message( event, message))
         api.connect()
 
+    def add_api(self, network_info):
+        self._create_api(mac_address=network_info['mac'], ip=network_info['host'], port=network_info['port'])
+
+
     async def handle_message(self, event: str, message: dict):
         mac = message['mac']
         if mac not in self.plugs.keys():
@@ -79,18 +146,55 @@ async def disconnect(self):
             if unsubscribe is not None:
                 unsubscribe()
 
+        await self.stop_processing_plug_queue()
+
     @callback
     def _acknowledge_sensor_added_to_homeassistant(self,mac, role):
         self.sensors[mac] = role
 
     @callback
-    def _plug_added(self, info):
-        _LOGGER.error(f" Request to add plug received: {info}")
+    async def _acknowledge_plug_added_to_homeassistant(self, mac_address, host, port):
+        self._create_api(mac_address, host, port)
+        await self._plug_added_queue.remove((mac_address, host, port))
+
+    @callback
+    async def _plug_added(self, info):
+        _LOGGER.debug(f" Request to add plug received: {info}")
+        network_info = dict()
+        network_info['mac'] = info['properties'][b'id'].decode('utf-8')
+        network_info['host'] = info['addresses'][0]
+        network_info['port'] = info['port']
+
+        if self._safe_to_process_plug_queue:
+            await self.enqueue_plug_for_adding(network_info)
+            await self.process_plug_queue()
+        else:
+            await self.enqueue_plug_for_adding(network_info)
 
     @callback
-    def _plug_updated(self, info):
-        _LOGGER.error(f" Request to update plug received: {info}")
+    async def _plug_updated(self, info):
+        _LOGGER.debug(f" Request to update plug received: {info}")
+        mac = info['properties'][b'id'].decode('utf-8')
+        host = info['addresses'][0]
+        port = info['port']
+        if mac in self.plugs:
+            self.plugs[mac].disconnect()
+
+        if mac in self._known_plugs:
+            self._create_api(mac, host, port)
+        else:
+            network_info = dict()
+            network_info['mac'] = mac
+            network_info['host'] = host
+            network_info['port'] = port
+            await self.enqueue_plug_for_adding(network_info)
+            await self.process_plug_queue()
 
     @callback
     def _plug_remove(self,name, info):
-        _LOGGER.error(f" Request to delete plug received: {info}")
+        _LOGGER.debug(f" Request to delete plug received: {info}")
+        mac = info['properties'][b'id'].decode('utf-8')
+        if mac in self.plugs:
+            self.plugs[mac].disconnect()
+
+        del self.plugs[mac]