Merge pull request #3 from DiUS/feature/zeroconf-discovery-background-task

All the hard work to resolve the issues around device discovery and disappearance.

Author: jmattsson

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/PowersensorDiscoveryService.py

@@ -0,0 +1,163 @@
+import asyncio
+from typing import Optional
+
+from homeassistant.core import HomeAssistant
+from zeroconf import ServiceBrowser, ServiceListener, Zeroconf
+from homeassistant.helpers.dispatcher import async_dispatcher_send
+import homeassistant.components.zeroconf
+
+from .const import DOMAIN
+
+import logging
+_LOGGER = logging.getLogger(__name__)
+
+
+class PowersensorServiceListener(ServiceListener):
+    def __init__(self, hass: HomeAssistant, debounce_timeout: float = 60):
+        self._hass = hass
+        self._plugs = {}
+        self._discoveries = {}
+        self._pending_removals = {}
+        self._debounce_seconds = debounce_timeout
+
+    def add_service(self, zc, type_, name):
+        self.cancel_any_pending_removal(name, "request to add")
+        info = self.__add_plug(zc, type_, name)
+        if info:
+            asyncio.run_coroutine_threadsafe(
+                self._async_service_add(self._plugs[name]),
+                self._hass.loop
+            )
+
+    async def _async_service_add(self, *args):
+        async_dispatcher_send(self._hass, f"{DOMAIN}_zeroconf_add_plug", *args)
+
+
+    async def _async_delayed_remove(self, name):
+        """Actually process the removal after delay."""
+        try:
+            await asyncio.sleep(self._debounce_seconds)
+            _LOGGER.info(f"Request to remove service {name} still pending after timeout. Processing remove request...")
+            if name in self._plugs:
+                data = self._plugs[name].copy()
+                del self._plugs[name]
+            else:
+                data = None
+            asyncio.run_coroutine_threadsafe(
+                self._async_service_remove(name, data),
+                self._hass.loop
+            )
+        except asyncio.CancelledError:
+            # Task was cancelled because service came back
+            _LOGGER.info(f"Request to remove service {name} was canceled by request to update or add plug.")
+            raise
+        finally:
+            # Either way were done with this task
+            self._pending_removals.pop(name, None)
+
+
+    def remove_service(self, zc, type_, name):
+        if name in self._pending_removals:
+            # removal for this service is already pending
+            return
+
+        _LOGGER.info(f"Scheduling removal for {name}")
+        self._pending_removals[name] = asyncio.run_coroutine_threadsafe(
+            self._async_delayed_remove(name),
+            self._hass.loop
+        )
+
+    async def _async_service_remove(self, *args):
+        async_dispatcher_send(self._hass, f"{DOMAIN}_zeroconf_remove_plug", *args)
+
+    def update_service(self, zc, type_, name):
+        self.cancel_any_pending_removal(name, "request to update")
+        info = self.__add_plug(zc, type_, name)
+        if info:
+            asyncio.run_coroutine_threadsafe(
+                self._async_service_update( self._plugs[name]),
+                self._hass.loop
+            )
+
+    async def _async_service_update(self, *args):
+        # remove from pending tasks if update received
+        async_dispatcher_send(self._hass, f"{DOMAIN}_zeroconf_update_plug", *args)
+
+    async def _async_get_service_info(self, zc, type_, name):
+        try:
+            info = await zc.async_get_service_info(type_, name, timeout=3000)
+            self._discoveries[name] = info
+        except Exception as e:
+            _LOGGER.error(f"Error retrieving info for {name}")
+
+
+    def __add_plug(self, zc, type_, name):
+        info = zc.get_service_info(type_, name)
+
+        if info:
+            self._plugs[name] = {'type': type_,
+                                 'name': name,
+                                 'addresses': ['.'.join(str(b) for b in addr) for addr in info.addresses],
+                                 'port': info.port,
+                                 'server': info.server,
+                                 'properties': info.properties
+                                 }
+        return info
+
+    def cancel_any_pending_removal(self, name, source):
+        task = self._pending_removals.pop(name, None)
+        if task:
+            task.cancel()
+            _LOGGER.info(f"Cancelled pending removal for {name} by {source}.")
+
+class PowersensorDiscoveryService:
+    def __init__(self, hass: HomeAssistant, service_type: str = "_powersensor._tcp.local."):
+        self._hass = hass
+        self.service_type = service_type
+
+        self.zc: Optional[Zeroconf] = None
+        self.listener: Optional[PowersensorServiceListener] = None
+        self.browser: Optional[ServiceBrowser] = None
+        self.running = False
+        self._task: Optional[asyncio.Task] = None
+
+    async def start(self):
+        """Start the mDNS discovery service"""
+        if self.running:
+            return
+
+        self.running = True
+        self.zc = await homeassistant.components.zeroconf.async_get_instance(self._hass)
+        self.listener = PowersensorServiceListener(self._hass)
+
+        # Create browser
+        self.browser = ServiceBrowser(self.zc, self.service_type, self.listener)
+
+        # Start the background task
+        self._task = asyncio.create_task(self._run())
+
+    async def _run(self):
+        """Background task that keeps the service alive"""
+        try:
+            while self.running:
+                await asyncio.sleep(1)
+        except asyncio.CancelledError:
+            pass
+
+    async def stop(self):
+        """Stop the mDNS discovery service"""
+        self.running = False
+
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+
+        if self.zc:
+            # self.zc.close()
+            self.zc = None
+
+        self.browser = None
+        self.listener = None

custom_components/powersensor/PowersensorEntity.py

@@ -0,0 +1,131 @@
+from datetime import timedelta
+from typing import Callable, Union
+
+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,
+                 input_config: dict[Union[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'Powersensor Device (ID: {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 = input_config[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."""
+
+        # event is not presently used, but is passed to maintain flexibility for future development
+
+        name_updated = False
+        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/PowersensorHouseholdEntity.py

@@ -22,6 +22,19 @@ class HouseholdMeasurements(Enum):
     ENERGY_TO_GRID = 7
     ENERGY_SOLAR_GENERATION = 8
 
+ConsumptionMeasurements = [
+    HouseholdMeasurements.POWER_HOME_USE,
+    HouseholdMeasurements.POWER_FROM_GRID,
+    HouseholdMeasurements.ENERGY_HOME_USE,
+    HouseholdMeasurements.ENERGY_FROM_GRID
+]
+ProductionMeasurements = [
+    HouseholdMeasurements.POWER_TO_GRID,
+    HouseholdMeasurements.POWER_SOLAR_GENERATION,
+    HouseholdMeasurements.ENERGY_TO_GRID,
+    HouseholdMeasurements.ENERGY_SOLAR_GENERATION
+]
+
 @dataclass
 class EntityConfig:
     name : str