trying this...again...

Author: bookman-dius

custom_components/powersensor/PowersensorDiscoveryService.py

@@ -1,6 +1,5 @@
 import asyncio
 from typing import Optional
-import logging
 
 from homeassistant.core import HomeAssistant
 from zeroconf import ServiceBrowser, ServiceListener, Zeroconf
@@ -9,6 +8,7 @@
 
 from .const import DOMAIN
 
+import logging
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -20,6 +20,7 @@ def __init__(self, hass: HomeAssistant, debounce_timeout: float = 60):
         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(
@@ -30,6 +31,7 @@ def add_service(self, zc, type_, name):
     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:
@@ -57,29 +59,30 @@ def remove_service(self, zc, type_, name):
             return
 
         _LOGGER.info(f"Scheduling removal for {name}")
-        self._pending_removals[name] = asyncio.create_task(self._async_delayed_remove(name))
-
+        self._pending_removals[name] = asyncio.run_coroutine_threadsafe(
+            self._async_delayed_remove(name),
+            self._hass.loop
+        )
+        print(self._pending_removals)
 
     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._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)
 
     def __add_plug(self, zc, type_, name):
         info = zc.get_service_info(type_, name)
-        task = self._pending_removals.pop(name, None)
-        if task:
-            task.cancel()
-            _LOGGER.info(f"Cancelled pending removal for {name}")
         if info:
             self._plugs[name] = {'type': type_,
                                  'name': name,
@@ -90,6 +93,11 @@ def __add_plug(self, zc, type_, name):
                                  }
         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."):

custom_components/powersensor/PowersensorHouseholdEntity.py

@@ -34,6 +34,7 @@ class HouseholdMeasurements(Enum):
     HouseholdMeasurements.ENERGY_TO_GRID,
     HouseholdMeasurements.ENERGY_SOLAR_GENERATION
 ]
+
 @dataclass
 class EntityConfig:
     name : str

custom_components/powersensor/PowersensorMessageDispatacher.py

@@ -1,4 +1,5 @@
 import asyncio
+import datetime
 import logging
 
 from homeassistant.core import HomeAssistant, callback
@@ -7,11 +8,11 @@
 from powersensor_local import PlugApi, VirtualHousehold
 
 from custom_components.powersensor.AsyncSet import AsyncSet
-from custom_components.powersensor.const import POWER_SENSOR_UPDATE_SIGNAL, DOMAIN, DEFAULT_PORT
+from custom_components.powersensor.const import POWER_SENSOR_UPDATE_SIGNAL, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 class PowersensorMessageDispatcher:
-    def __init__(self, hass: HomeAssistant, vhh: VirtualHousehold,debounce_timeout: float = 60):
+    def __init__(self, hass: HomeAssistant, vhh: VirtualHousehold, debounce_timeout: float = 60):
         self._hass = hass
         self._vhh = vhh
         self.plugs = dict()
@@ -21,6 +22,11 @@ def __init__(self, hass: HomeAssistant, vhh: VirtualHousehold,debounce_timeout:
         self.on_start_sensor_queue = dict()
         self._pending_removals = {}
         self._debounce_seconds = debounce_timeout
+        self.has_solar = False
+        self._virtual_household_has_been_setup = False
+        self._last_request_to_notify_about_solar = datetime.datetime(1970,1,1,0,0,0)
+        self._solar_request_limit = datetime.timedelta(seconds = 10)
+
         self._unsubscribe_from_signals = [
             async_dispatcher_connect(self._hass,
                                      f"{DOMAIN}_sensor_added_to_homeassistant",
@@ -37,6 +43,9 @@ def __init__(self, hass: HomeAssistant, vhh: VirtualHousehold,debounce_timeout:
             async_dispatcher_connect(self._hass,
                                      f"{DOMAIN}_plug_added_to_homeassistant",
                                      self._acknowledge_plug_added_to_homeassistant),
+            async_dispatcher_connect(self._hass,
+                                     f"{DOMAIN}_solar_added_to_virtual_household",
+                                     self._acknowledge_solar_added_to_virtual_household),
         ]
 
         self._monitor_add_plug_queue = None
@@ -98,6 +107,21 @@ async def stop_processing_plug_queue(self):
             _LOGGER.debug("Background task stopped")
             self._monitor_add_plug_queue = None
 
+
+    async def stop_pending_removal_tasks(self):
+        """Stop the background removal tasks."""
+        for k in range(len(self._pending_removals)):
+            if self._pending_removals[k] and not self._pending_removals[k].done():
+                self._pending_removals[k].cancel()
+                try:
+                    await self._pending_removals[k]
+                except asyncio.CancelledError:
+                    pass
+                _LOGGER.debug("Background removal task stopped")
+                self._pending_removals[k] = None
+        self._pending_removals = []
+
+
     def _create_api(self, mac_address, ip, port, name):
         _LOGGER.info(f"Creating API for mac={mac_address}, ip={ip}, port={port}")
         api = PlugApi(mac=mac_address, ip=ip, port=port)
@@ -126,28 +150,36 @@ def add_api(self, network_info):
         self._create_api(mac_address=network_info['mac'], ip=network_info['host'],
                          port=network_info['port'], name=network_info['name'])
 
-
-    async def handle_message(self, event: str, message: dict):
-        mac = message['mac']
+    def cancel_any_pending_removal(self, mac, source):
         task = self._pending_removals.pop(mac, None)
         if task:
             task.cancel()
-            _LOGGER.info(f"Cancelled pending removal for {mac} by new message received from plug.")
+            _LOGGER.info(f"Cancelled pending removal for {mac} by {source}.")
+
+    async def handle_message(self, event: str, message: dict):
+        mac = message['mac']
+        role = message.get('role', None)
+        self.cancel_any_pending_removal(mac, "new message received from plug")
 
         if mac not in self.plugs.keys():
             if mac not in self.sensors:
-                role = None
-                if 'role' in message:
+                if role is not None:
                     self.on_start_sensor_queue[mac] = role
-                    role = message['role']
                 async_dispatcher_send(self._hass, f"{DOMAIN}_create_sensor", mac, role)
 
         # Feed the household calculations
         if event == 'average_power':
             await self._vhh.process_average_power_event(message)
         elif event == 'summation_energy':
             await self._vhh.process_summation_event(message)
-
+        if role == 'solar':
+            self.has_solar = True
+            if not self._virtual_household_has_been_setup:
+                new_time = datetime.datetime.now()
+                if self._last_request_to_notify_about_solar + self._solar_request_limit <new_time:
+                    self._last_request_to_notify_about_solar = new_time
+                    _LOGGER.debug("Notifying integration that solar is present.")
+                    async_dispatcher_send(self._hass, f"{DOMAIN}_solar_sensor_detected")
         async_dispatcher_send(self._hass, f"{POWER_SENSOR_UPDATE_SIGNAL}_{mac}_{event}", event, message)
 
     async def disconnect(self):
@@ -159,11 +191,17 @@ async def disconnect(self):
                 unsubscribe()
 
         await self.stop_processing_plug_queue()
+        await self.stop_pending_removal_tasks()
 
     @callback
     def _acknowledge_sensor_added_to_homeassistant(self,mac, role):
         self.sensors[mac] = role
 
+    @callback
+    def _acknowledge_solar_added_to_virtual_household(self, success):
+        _LOGGER.debug("Solar has been added to virtual household.")
+        self._virtual_household_has_been_setup = success
+
     async def _acknowledge_plug_added_to_homeassistant(self, mac_address, host, port, name):
         _LOGGER.info(f"Adding new API for mac={mac_address}, ip={host}, port={port}")
         self._create_api(mac_address, host, port, name)
@@ -174,15 +212,11 @@ async def _plug_added(self, info):
         network_info = dict()
         mac = info['properties'][b'id'].decode('utf-8')
         network_info['mac'] = mac
+        self.cancel_any_pending_removal(mac, "request to add plug")
         network_info['host'] = info['addresses'][0]
         network_info['port'] = info['port']
         network_info['name'] = info['name']
 
-        task = self._pending_removals.pop(mac, None)
-        if task:
-            task.cancel()
-            _LOGGER.info(f"Cancelled pending removal for {mac} by request to add api.")
-
         if self._safe_to_process_plug_queue:
             await self.enqueue_plug_for_adding(network_info)
             await self.process_plug_queue()
@@ -192,13 +226,11 @@ async def _plug_added(self, info):
     async def _plug_updated(self, info):
         _LOGGER.debug(f" Request to update plug received: {info}")
         mac = info['properties'][b'id'].decode('utf-8')
+        self.cancel_any_pending_removal(mac, "request to update plug")
         host = info['addresses'][0]
         port = info['port']
         name = info['name']
-        task = self._pending_removals.pop(mac, None)
-        if task:
-            task.cancel()
-            _LOGGER.info(f"Cancelled pending removal for {mac} by request to update api.")
+
         if mac in self.plugs:
             current_api: PlugApi = self.plugs[mac]
             if current_api._listener._ip == host and current_api._listener._port == port:
@@ -218,6 +250,7 @@ async def _plug_updated(self, info):
             await self.enqueue_plug_for_adding(network_info)
             await self.process_plug_queue()
 
+
     async def _schedule_plug_removal(self, name, info):
         _LOGGER.debug(f" Request to delete plug received: {info}")
         if name in self._known_plug_names:
@@ -228,7 +261,10 @@ async def _schedule_plug_removal(self, name, info):
                     return
 
                 _LOGGER.info(f"Scheduling removal for {name}")
-                self._pending_removals[name] = asyncio.create_task(self._delayed_plug_remove(name,mac))
+                self._pending_removals[mac] = self._hass.async_create_background_task(
+                    self._delayed_plug_remove(name,mac),
+                    name = f"Removal-Task-For-{name}"
+                )
         else:
             _LOGGER.warning(f"Received request to delete api for gateway with name [{name}], but this name"
                             f"is not associated with an existing PlugAPI. Ignoring...")
@@ -242,8 +278,8 @@ async def _delayed_plug_remove(self, name, mac):
             del self.plugs[mac]
             del self._known_plug_names[name]
         except asyncio.CancelledError:
-            # Task was cancelled because service came back
-            _LOGGER.info(f"Request to remove plug {mac} was canceled by request to update, add plug or new message.")
+            # Task was canceled because service came back
+            _LOGGER.info(f"Request to remove plug {mac} was cancelled by request to update, add plug or new message.")
 
         # Either way were done with this task
-        self._pending_removals.pop(name, None)
+        self._pending_removals.pop(name, None)

custom_components/powersensor/__init__.py

@@ -21,24 +21,24 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     hass.data.setdefault(DOMAIN, {})
     hass.data[DOMAIN][entry.entry_id] = {}
-    my_data = hass.data[DOMAIN][entry.entry_id]
 
     integration = await async_get_integration(hass, DOMAIN)
     manifest  = integration.manifest
 
     # Establish create the zeroconf discovery service
-    zeroconf_listener = PowersensorDiscoveryService(hass, manifest["zeroconf"][0])
-    await zeroconf_listener.start()
+    zeroconf_service= PowersensorDiscoveryService(hass, manifest["zeroconf"][0])
+    await zeroconf_service.start()
 
     # Establish our virtual household
-    vhh = VirtualHousehold(my_data["with_solar"] if "with_solar" in my_data else False)
-    
+    vhh = VirtualHousehold(False)
+
 
     # TODO: can we move the dispatcher into the entry.runtime_data dict?
     dispatcher = PowersensorMessageDispatcher(hass, vhh)
     for mac, network_info in entry.data.items():
         await dispatcher.enqueue_plug_for_adding(network_info)
-    entry.runtime_data = {"vhh": vhh, "zeroconf": zeroconf_listener, "dispatcher": dispatcher}
+
+    entry.runtime_data = { "vhh": vhh , "dispatcher" : dispatcher, "zeroconf" : zeroconf_service}
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
 
@@ -47,11 +47,10 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     _LOGGER.debug("Started unloading for %s", entry.entry_id)
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        if entry.runtime_data:
-            if "dispatcher" in entry.runtime_data.keys():
-                    await entry.runtime_data["dispatcher"].disconnect()
-            if "zeroconf" in entry.runtime_data.keys():
-                await entry.runtime_data["zeroconf"].stop()
+        if "dispatcher" in entry.runtime_data .keys():
+            await entry.runtime_data["dispatcher"].disconnect()
+        if "zeroconf" in entry.runtime_data .keys():
+            await entry.runtime_data["zeroconf"].stop()
 
     return unload_ok
 

custom_components/powersensor/sensor.py

@@ -10,7 +10,8 @@
 
 from . import PowersensorMessageDispatcher
 from .PlugMeasurements import PlugMeasurements
-from .PowersensorHouseholdEntity import ConsumptionMeasurements, ProductionMeasurements, PowersensorHouseholdEntity
+from .PowersensorHouseholdEntity import HouseholdMeasurements, PowersensorHouseholdEntity, ConsumptionMeasurements, \
+    ProductionMeasurements
 from .PowersensorPlugEntity import PowersensorPlugEntity
 from .PowersensorSensorEntity import PowersensorSensorEntity
 from .SensorMeasurements import SensorMeasurements
@@ -25,8 +26,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Powersensor sensors."""
     vhh = entry.runtime_data["vhh"]
-    my_data = hass.data[DOMAIN][entry.entry_id]
-    dispatcher: PowersensorMessageDispatcher =  entry.runtime_data['dispatcher']
+    dispatcher: PowersensorMessageDispatcher = entry.runtime_data['dispatcher']
 
 
     async def create_plug(plug_mac_address: str):
@@ -54,7 +54,7 @@ async def handle_discovered_plug(plug_mac_address: str, host: str, port: int, na
 
     async def handle_discovered_sensor(sensor_mac: str, sensor_role: str):
         if sensor_role == 'solar':
-            my_data["with_solar"] = True  # Remember for next time we start
+            entry.runtime_data["with_solar"] = True  # Remember for next time we start
 
         new_sensors = [
             PowersensorSensorEntity(hass, sensor_mac, SensorMeasurements.Battery),
@@ -76,15 +76,24 @@ async def handle_discovered_sensor(sensor_mac: str, sensor_role: str):
     for mac, role in dispatcher.on_start_sensor_queue.items():
         await handle_discovered_sensor(mac, role)
 
+    async def add_solar_to_virtual_household():
+        _LOGGER.debug("Received request to add solar to virtual household")
+        solar_household_entities = []
+        for solar_measurement_type in ProductionMeasurements:
+            solar_household_entities.append(PowersensorHouseholdEntity(vhh, solar_measurement_type))
+
+        async_add_entities(solar_household_entities)
+        async_dispatcher_send(hass, f"{DOMAIN}_solar_added_to_virtual_household", True)
+
+    entry.async_on_unload(
+        async_dispatcher_connect(
+            hass, f"{DOMAIN}_solar_sensor_detected", add_solar_to_virtual_household
+        )
+    )
     # Register the virtual household entities
     household_entities = []
     for measurement_type in ConsumptionMeasurements:
         household_entities.append(PowersensorHouseholdEntity(vhh, measurement_type))
-
-
-    for measurement_type in ProductionMeasurements:
-        household_entities.append(PowersensorHouseholdEntity(vhh, measurement_type))
-
     async_add_entities(household_entities)
 
     async_dispatcher_send(hass, f"{DOMAIN}_setup_complete", True)