Fix handling disconnected devices

custom_components/googlehome/__init__.py

@@ -76,31 +76,33 @@ def __init__(self, hass):
         """Initialize the Google Home Client."""
         self.hass = hass
 
-    async def update_info(self, host):
+    async def update_info(self, host, uuid):
         """Update data from Google Home."""
         from googledevices.api.connect import Cast
 
         _LOGGER.debug("Updating Google Home info for %s", host)
         session = async_get_clientsession(self.hass)
 
+        # eureka_info might not return something, if the device just started up
+        await asyncio.sleep(10)
         device_info = await Cast(host, self.hass.loop, session).info()
         for token in self.hass.data[TOKENS].values():
             device_info_data = await device_info.get_device_info(token)
             if device_info_data is not None:
                 _LOGGER.debug(device_info_data)
-                self.hass.data[DOMAIN][host]["info"] = device_info_data
+                self.hass.data[DOMAIN][uuid]["info"] = device_info_data
                 return True
         return False
 
-    async def update_bluetooth(self, host, entry: config_entries.ConfigEntry):
+    async def update_bluetooth(self, host, uuid, entry: config_entries.ConfigEntry):
         """Update bluetooth from Google Home."""
         from googledevices.api.connect import Cast
 
         _LOGGER.debug("Updating Google Home bluetooth for %s", host)
         session = async_get_clientsession(self.hass)
 
         try:
-            info = self.hass.data[DOMAIN][host]["info"]
+            info = self.hass.data[DOMAIN][uuid]["info"]
             token = self.hass.data[TOKENS][info["device_info"]["cloud_device_id"]]
         except KeyError:
             return
@@ -115,21 +117,21 @@ async def update_bluetooth(self, host, entry: config_entries.ConfigEntry):
             return
 
         _LOGGER.debug(bluetooth_data)
-        self.hass.data[DOMAIN][host]["bluetooth"] = bluetooth_data
+        self.hass.data[DOMAIN][uuid]["bluetooth"] = bluetooth_data
 
-    async def update_alarms(self, host, entry: config_entries.ConfigEntry):
+    async def update_alarms(self, host, uuid, entry: config_entries.ConfigEntry):
         """Update alarms from Google Home."""
         from googledevices.api.connect import Cast
 
         _LOGGER.debug("Updating Google Home alarm for %s", host)
         session = async_get_clientsession(self.hass)
 
         try:
-            info = self.hass.data[DOMAIN][host]["info"]
+            info = self.hass.data[DOMAIN][uuid]["info"]
             token = self.hass.data[TOKENS][info["device_info"]["cloud_device_id"]]
         except KeyError:
-            if not self.hass.data[DOMAIN][host].get("alarms"):
-                self.hass.data[DOMAIN][host]["alarms"] = {"timer": [], "alarm": []}
+            if not self.hass.data[DOMAIN][uuid].get("alarms"):
+                self.hass.data[DOMAIN][uuid]["alarms"] = {"timer": [], "alarm": []}
             return
 
         assistant = await Cast(host, self.hass.loop, session).assistant()
@@ -140,4 +142,4 @@ async def update_alarms(self, host, entry: config_entries.ConfigEntry):
             return
 
         _LOGGER.debug(alarms_data)
-        self.hass.data[DOMAIN][host]["alarms"] = alarms_data
+        self.hass.data[DOMAIN][uuid]["alarms"] = alarms_data

custom_components/googlehome/device_tracker.py

@@ -1,6 +1,10 @@
 """Support for Google Home Bluetooth tacker."""
 from datetime import timedelta
 import logging
+from pychromecast.socket_client import (
+    CONNECTION_STATUS_CONNECTED,
+    CONNECTION_STATUS_DISCONNECTED,
+)
 
 from homeassistant.components.device_tracker import DeviceScanner
 from homeassistant.helpers.event import async_track_time_interval
@@ -16,38 +20,42 @@
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
 from .const import CLIENT, DOMAIN, NAME, CONF_RSSI_THRESHOLD, CONF_DEVICE_TYPES, DEFAULT_RSSI_THRESHOLD, DEFAULT_DEVICE_TYPES
+from .helpers import ChromecastMonitor
+
 
 _LOGGER = logging.getLogger(__name__)
 
 DEFAULT_SCAN_INTERVAL = timedelta(seconds=10)
 
-active_devices = {}
-
+monitors = []
 async def async_setup_entry(hass, config_entry, async_see):
     """Setup the Google Home scanner platform"""
-    async def async_cast_discovered(discover: ChromecastInfo):
-        if discover.is_audio_group:
-            return
-        hass.data[DOMAIN].setdefault(discover.host, {})
-
-        if await hass.data[CLIENT].update_info(discover.host):
-            info = hass.data[DOMAIN][discover.host]["info"]
-            if info["device_info"]["cloud_device_id"] not in active_devices and info["device_info"]["capabilities"].get("bluetooth_supported", False):
-                scanner = GoogleHomeDeviceScanner(
-                    hass, hass.data[CLIENT], config_entry, discover, async_see
-                )
-                active_devices[info["device_info"]["cloud_device_id"]] = scanner
-                await scanner.async_init()
-
-    async def async_cast_removed(info: ChromecastInfo):
-        if info["device_info"]["cloud_device_id"] in active_devices:
-            await active_devices[info["device_info"]["cloud_device_id"]].async_deinit()
-            del active_devices[info["device_info"]["cloud_device_id"]]
-
-    async_dispatcher_connect(hass, SIGNAL_CAST_REMOVED, async_cast_removed)
-    async_dispatcher_connect(hass, SIGNAL_CAST_DISCOVERED, async_cast_discovered)
-    for chromecast in hass.data[KNOWN_CHROMECAST_INFO_KEY].values():
-        await async_cast_discovered(chromecast)
+    monitor = DeviceTrackerMonitor()
+    await monitor.async_init(hass, config_entry, async_see)
+    monitors.append(monitor)
+
+class DeviceTrackerMonitor(ChromecastMonitor):
+    async def async_init(self, hass, config_entry, async_see):
+        self._config_entry = config_entry
+        self._async_see = async_see
+        await super().async_init(hass)
+
+    async def supported(self, discover: ChromecastInfo) -> bool:
+        info = self._hass.data[DOMAIN][discover.uuid]["info"]
+        return info["device_info"]["capabilities"].get("bluetooth_supported", False)
+
+    async def setup(self, discover: ChromecastInfo):
+        scanner = GoogleHomeDeviceScanner(
+            self._hass, self._hass.data[CLIENT], self._config_entry, discover, self._async_see
+        )
+        await scanner.async_init()
+        return [scanner]
+
+    async def cleanup(self, discover):
+        info = self._hass.data[DOMAIN][discover]["info"]
+        for entity in self._active_devices[info["device_info"]["cloud_device_id"]]:
+            await entity.async_deinit()
+
 
 class GoogleHomeDeviceScanner():
     """This class queries a Google Home unit."""
@@ -59,30 +67,46 @@ def __init__(self, hass, client, config, device, async_see):
         self.rssi = config.options.get(CONF_RSSI_THRESHOLD, DEFAULT_RSSI_THRESHOLD),
         self.device_types = config.options.get(CONF_DEVICE_TYPES, DEFAULT_DEVICE_TYPES),
         self.host = device.host
+        self.uuid = device.uuid
         self.name = device.friendly_name
         self.client = client
         self.config_entry = config
         self.removal = None
+        self.active = True
+
+    def new_connection_status(self, connection_status):
+        if connection_status == CONNECTION_STATUS_CONNECTED:
+            self.active = True
+        elif connection_status == CONNECTION_STATUS_DISCONNECTED:
+            self.active = False
 
     async def async_init(self):
         """Further initialize connection to Google Home."""
-        data = self.hass.data[DOMAIN][self.host]
+        data = self.hass.data[DOMAIN][self.uuid]
         info = data.get("info", {})
         await self.async_update()
         self.removal = async_track_time_interval(
             self.hass, self.async_update, DEFAULT_SCAN_INTERVAL
         )
+        self.active = True
 
     async def async_deinit(self):
         if self.removal:
             self.removal()
         self.removal = None
+        self.active = False
+
+    async def async_remove(self):
+        self.async_deinit()
 
     async def async_update(self, now=None):
         """Ensure the information from Google Home is up to date."""
+        if not self.active:
+            return
+
         _LOGGER.debug("Checking Devices on %s", self.host)
-        await self.client.update_bluetooth(self.host, self.config_entry)
-        data = self.hass.data[DOMAIN][self.host]
+        await self.client.update_bluetooth(self.host, self.uuid, self.config_entry)
+        data = self.hass.data[DOMAIN][self.uuid]
         info = data.get("info")
         bluetooth = data.get("bluetooth")
         if info is None or bluetooth is None:

custom_components/googlehome/helpers.py

@@ -0,0 +1,72 @@
+from homeassistant.components.cast.const import (
+    SIGNAL_CAST_DISCOVERED,
+    SIGNAL_CAST_REMOVED,
+    KNOWN_CHROMECAST_INFO_KEY,
+    DOMAIN as CAST_DOMAIN,
+)
+from homeassistant.components.cast.helpers import ChromecastInfo, ChromeCastZeroconf
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
+
+import pychromecast
+import logging
+
+from .const import CLIENT, DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+class ChromecastMonitor:
+    async def async_init(self, hass):
+        self._hass = hass
+        self._active_devices = {}
+
+        _LOGGER.debug("Setting up signals")
+        async_dispatcher_connect(hass, SIGNAL_CAST_REMOVED, self.async_cast_removed)
+        async_dispatcher_connect(hass, SIGNAL_CAST_DISCOVERED, self.async_cast_discovered)
+        for chromecast in hass.data[KNOWN_CHROMECAST_INFO_KEY].copy().values():
+            await self.async_cast_discovered(chromecast)
+
+    async def async_cast_discovered(self, discover: ChromecastInfo):
+        _LOGGER.debug("Discovered {}".format(discover.host))
+        if discover.is_audio_group:
+            return
+
+        self._hass.data[DOMAIN].setdefault(discover.uuid, {})
+        if await self._hass.data[CLIENT].update_info(discover.host, discover.uuid):
+            info = self._hass.data[DOMAIN][discover.uuid]["info"]
+            if info["device_info"]["cloud_device_id"] not in self._active_devices and await self.supported(discover):
+                chromecast = await self._hass.async_add_executor_job(
+                    pychromecast.get_chromecast_from_service,
+                    (
+                        discover.services,
+                        discover.uuid,
+                        discover.model_name,
+                        discover.friendly_name,
+                        None,
+                        None,
+                    ),
+                    ChromeCastZeroconf.get_zeroconf(),
+                )
+                chromecast.register_connection_listener(self)
+                self._active_devices[info["device_info"]["cloud_device_id"]] = await self.setup(discover)
+
+    async def async_cast_removed(self, discover: ChromecastInfo):
+        _LOGGER.debug("Removed {}".format(discover.uuid))
+
+        info = self._hass.data[DOMAIN][discover.uuid]["info"]
+        if info["device_info"]["cloud_device_id"] in self._active_devices:
+            await self.cleanup(discover.uuid)
+            del self._active_devices[info["device_info"]["cloud_device_id"]]
+
+    async def supported(self, discover: ChromecastInfo) -> bool:
+        return True
+
+    async def setup(self, discover: ChromecastInfo):
+        return []
+
+    async def cleanup(self, discover: ChromecastInfo):
+        pass
+
+    def new_connection_status(self, connection_status):
+        """Handle reception of a new ConnectionStatus."""
+        for entity in self._active_devices.values():
+            entity.connection_status(connection_status.status)