Merge pull request #15 from djtimca/feature-dynamicentities

Enable auto-adding of new entities (switch, device_tracker) when a new device connects to the network.

custom_components/googlewifi/__init__.py

@@ -11,6 +11,7 @@
 from homeassistant.core import CoreState, HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady, PlatformNotReady
 from homeassistant.helpers import aiohttp_client
+from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
     DataUpdateCoordinator,
@@ -28,6 +29,8 @@
     GOOGLEWIFI_API,
     POLLING_INTERVAL,
     REFRESH_TOKEN,
+    SIGNAL_ADD_DEVICE,
+    SIGNAL_DELETE_DEVICE,
 )
 
 CONFIG_SCHEMA = vol.Schema({DOMAIN: vol.Schema({})}, extra=vol.ALLOW_EXTRA)
@@ -112,6 +115,17 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
     return unload_ok
 
 
+async def cleanup_device_registry(hass: HomeAssistant, device_id):
+    """Remove device registry entry if there are no remaining entities."""
+
+    device_registry = await hass.helpers.device_registry.async_get_registry()
+    entity_registry = await hass.helpers.entity_registry.async_get_registry()
+    if device_id and not hass.helpers.entity_registry.async_entries_for_device(
+        entity_registry, device_id, include_disabled_entities=True
+    ):
+        device_registry.async_remove_device(device_id)
+
+
 class GoogleWiFiUpdater(DataUpdateCoordinator):
     """Class to manage fetching update data from the Google Wifi API."""
 
@@ -136,6 +150,7 @@ def __init__(
         self.auto_speedtest = auto_speedtest
         self.speedtest_interval = speedtest_interval
         self._force_speed_update = None
+        self.devicelist = []
 
         super().__init__(
             hass=hass,
@@ -167,6 +182,15 @@ async def _async_update_data(self):
                     device_network = device.get("ipAddress", " " * 10)
                     device_network = ".".join(device_network.split(".", 3)[:3])
 
+                    if device_id not in self.devicelist:
+                        to_add = {
+                            "system_id": system_id,
+                            "device_id": device_id,
+                            "device": device,
+                        }
+                        async_dispatcher_send(self.hass, SIGNAL_ADD_DEVICE, to_add)
+                        self.devicelist.append(device_id)
+
                     if device.get("connected") and main_network == device_network:
                         connected_count += 1
                         device["network"] = "main"
@@ -180,6 +204,13 @@ async def _async_update_data(self):
                         connected_count += 1
                         device["network"] = "main"
 
+                for known_device in self.devicelist:
+                    if known_device not in system["devices"]:
+                        async_dispatcher_send(
+                            self.hass, SIGNAL_DELETE_DEVICE, known_device
+                        )
+                        self.devicelist.remove(known_device)
+
             system_data[system_id]["connected_devices"] = connected_count
             system_data[system_id]["guest_devices"] = guest_connected_count
             system_data[system_id]["total_devices"] = (
@@ -274,3 +305,18 @@ async def async_added_to_hass(self):
     def _update_callback(self):
         """Handle device update."""
         self.async_write_ha_state()
+
+    async def _delete_callback(self, device_id):
+        """Remove the device when it disappears."""
+
+        if device_id == self._unique_id:
+            entity_registry = (
+                await self.hass.helpers.entity_registry.async_get_registry()
+            )
+
+            if entity_registry.async_is_registered(self.entity_id):
+                entity_entry = entity_registry.async_get(self.entity_id)
+                entity_registry.async_remove(self.entity_id)
+                await cleanup_device_registry(self.hass, entity_entry.device_id)
+            else:
+                await self.async_remove()

custom_components/googlewifi/const.py

@@ -28,6 +28,8 @@
 CONF_SPEEDTEST_INTERVAL = "speedtest_interval"
 DEFAULT_SPEEDTEST_INTERVAL = 24
 CONF_SPEED_UNITS = "speed_units"
+SIGNAL_ADD_DEVICE = "googlewifi_add_device"
+SIGNAL_DELETE_DEVICE = "googlewifi_delete_device"
 
 
 def unit_convert(data_rate: float, unit_of_measurement: str):

custom_components/googlewifi/device_tracker.py

@@ -5,6 +5,7 @@
 from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER
 from homeassistant.const import ATTR_NAME
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
 from . import GoogleWifiEntity, GoogleWiFiUpdater
 from .const import (
@@ -17,6 +18,7 @@
     DEV_CLIENT_MODEL,
     DEV_MANUFACTURER,
     DOMAIN,
+    SIGNAL_ADD_DEVICE,
 )
 
 
@@ -44,6 +46,29 @@ async def async_setup_entry(hass, entry, async_add_entities):
 
     async_add_entities(entities)
 
+    async def async_new_entities(device_info):
+        """Add new entities when they connect to Google Wifi."""
+        system_id = device_info["system_id"]
+        device_id = device_info["device_id"]
+        device = device_info["device"]
+
+        device_name = f"{device['friendlyName']}"
+
+        if device.get("friendlyType"):
+            device_name = device_name + f" ({device['friendlyType']})"
+
+        entity = GoogleWifiDeviceTracker(
+            coordinator=coordinator,
+            name=device_name,
+            icon=DEFAULT_ICON,
+            system_id=system_id,
+            item_id=device_id,
+        )
+        entities = [entity]
+        async_add_entities(entities)
+
+    async_dispatcher_connect(hass, SIGNAL_ADD_DEVICE, async_new_entities)
+
 
 class GoogleWifiDeviceTracker(GoogleWifiEntity, ScannerEntity):
     """Defines a Google WiFi device tracker."""

custom_components/googlewifi/switch.py

@@ -4,9 +4,11 @@
 import voluptuous as vol
 from homeassistant.components.switch import SwitchEntity
 from homeassistant.const import ATTR_NAME, DATA_RATE_MEGABYTES_PER_SECOND
+from homeassistant.core import callback
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers import entity_platform
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.util.dt import as_local, as_timestamp, parse_datetime
 
 from . import GoogleWifiEntity, GoogleWiFiUpdater
@@ -21,6 +23,8 @@
     DEV_CLIENT_MODEL,
     DOMAIN,
     PAUSE_UPDATE,
+    SIGNAL_ADD_DEVICE,
+    SIGNAL_DELETE_DEVICE,
     unit_convert,
 )
 
@@ -32,6 +36,8 @@ async def async_setup_entry(hass, entry, async_add_entities):
     """Set up the switch platform."""
 
     coordinator = hass.data[DOMAIN][entry.entry_id][COORDINATOR]
+    device = hass.data[DOMAIN][entry.entry_id]
+
     entities = []
 
     data_unit = entry.options.get(CONF_SPEED_UNITS, DATA_RATE_MEGABYTES_PER_SECOND)
@@ -55,6 +61,30 @@ async def async_setup_entry(hass, entry, async_add_entities):
 
     async_add_entities(entities)
 
+    async def async_new_entities(device_info):
+        """Add new entities when they connect to Google Wifi."""
+        system_id = device_info["system_id"]
+        device_id = device_info["device_id"]
+        device = device_info["device"]
+
+        device_name = f"{device['friendlyName']}"
+
+        if device.get("friendlyType"):
+            device_name = device_name + f" ({device['friendlyType']})"
+
+        entity = GoogleWifiSwitch(
+            coordinator=coordinator,
+            name=device_name,
+            icon=DEFAULT_ICON,
+            system_id=system_id,
+            item_id=device_id,
+            data_unit=data_unit,
+        )
+        entities = [entity]
+        async_add_entities(entities)
+
+    async_dispatcher_connect(hass, SIGNAL_ADD_DEVICE, async_new_entities)
+
     # register service for reset
     platform = entity_platform.current_platform.get()
 
@@ -70,6 +100,8 @@ async def async_setup_entry(hass, entry, async_add_entities):
         "async_clear_prioritization",
     )
 
+    return True
+
 
 class GoogleWifiSwitch(GoogleWifiEntity, SwitchEntity):
     """Defines a Google WiFi switch."""