Cast unique_id and async discovery (#12474)

* Cast unique_id and async discovery

* Lazily load chromecasts

* Lint

* Fixes & Improvements

* Fixes

* Improve disconnects

cast.disconnect with blocking=False does **not** do I/O; it simply sets an event for the socket client looper

* Add tests

* Remove unnecessary calls

* Lint

* Fix use of hass object

homeassistant/components/media_player/cast.py

@@ -5,10 +5,16 @@ For more details about this platform, please refer to the documentation at
 https://home-assistant.io/components/media_player.cast/
 """
 # pylint: disable=import-error
+import asyncio
 import logging
+import threading
 
 import voluptuous as vol
 
+from homeassistant.helpers.typing import HomeAssistantType, ConfigType
+from homeassistant.core import callback
+from homeassistant.helpers.dispatcher import (dispatcher_send,
+                                              async_dispatcher_connect)
 from homeassistant.components.media_player import (
     MEDIA_TYPE_MUSIC, MEDIA_TYPE_TVSHOW, MEDIA_TYPE_VIDEO, SUPPORT_NEXT_TRACK,
     SUPPORT_PAUSE, SUPPORT_PLAY_MEDIA, SUPPORT_PREVIOUS_TRACK,
@@ -16,7 +22,7 @@ from homeassistant.components.media_player import (
     SUPPORT_STOP, SUPPORT_PLAY, MediaPlayerDevice, PLATFORM_SCHEMA)
 from homeassistant.const import (
     CONF_HOST, STATE_IDLE, STATE_OFF, STATE_PAUSED, STATE_PLAYING,
-    STATE_UNKNOWN)
+    STATE_UNKNOWN, EVENT_HOMEASSISTANT_STOP)
 import homeassistant.helpers.config_validation as cv
 import homeassistant.util.dt as dt_util
 
@@ -33,7 +39,13 @@ SUPPORT_CAST = SUPPORT_PAUSE | SUPPORT_VOLUME_SET | SUPPORT_VOLUME_MUTE | \
     SUPPORT_TURN_ON | SUPPORT_TURN_OFF | SUPPORT_PREVIOUS_TRACK | \
     SUPPORT_NEXT_TRACK | SUPPORT_PLAY_MEDIA | SUPPORT_STOP | SUPPORT_PLAY
 
-KNOWN_HOSTS_KEY = 'cast_known_hosts'
+INTERNAL_DISCOVERY_RUNNING_KEY = 'cast_discovery_running'
+# UUID -> CastDevice mapping; cast devices without UUID are not stored
+ADDED_CAST_DEVICES_KEY = 'cast_added_cast_devices'
+# Stores every discovered (host, port, uuid)
+KNOWN_CHROMECASTS_KEY = 'cast_all_chromecasts'
+
+SIGNAL_CAST_DISCOVERED = 'cast_discovered'
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Optional(CONF_HOST): cv.string,
@@ -41,67 +53,144 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
 })
 
 
-# pylint: disable=unused-argument
+def _setup_internal_discovery(hass: HomeAssistantType) -> None:
-def setup_platform(hass, config, add_devices, discovery_info=None):
+    """Set up the pychromecast internal discovery."""
+    hass.data.setdefault(INTERNAL_DISCOVERY_RUNNING_KEY, threading.Lock())
+    if not hass.data[INTERNAL_DISCOVERY_RUNNING_KEY].acquire(blocking=False):
+        # Internal discovery is already running
+        return
+
+    import pychromecast
+
+    def internal_callback(name):
+        """Called when zeroconf has discovered a new chromecast."""
+        mdns = listener.services[name]
+        ip_address, port, uuid, _, _ = mdns
+        key = (ip_address, port, uuid)
+
+        if key in hass.data[KNOWN_CHROMECASTS_KEY]:
+            _LOGGER.debug("Discovered previous chromecast %s", mdns)
+            return
+
+        _LOGGER.debug("Discovered new chromecast %s", mdns)
+        try:
+            # pylint: disable=protected-access
+            chromecast = pychromecast._get_chromecast_from_host(
+                mdns, blocking=True)
+        except pychromecast.ChromecastConnectionError:
+            _LOGGER.debug("Can't set up cast with mDNS info %s. "
+                          "Assuming it's not a Chromecast", mdns)
+            return
+        hass.data[KNOWN_CHROMECASTS_KEY][key] = chromecast
+        dispatcher_send(hass, SIGNAL_CAST_DISCOVERED, chromecast)
+
+    _LOGGER.debug("Starting internal pychromecast discovery.")
+    listener, browser = pychromecast.start_discovery(internal_callback)
+
+    def stop_discovery(event):
+        """Stop discovery of new chromecasts."""
+        pychromecast.stop_discovery(browser)
+        hass.data[INTERNAL_DISCOVERY_RUNNING_KEY].release()
+
+    hass.bus.listen_once(EVENT_HOMEASSISTANT_STOP, stop_discovery)
+
+
+@callback
+def _async_create_cast_device(hass, chromecast):
+    """Create a CastDevice Entity from the chromecast object.
+
+    Returns None if the cast device has already been added. Additionally,
+    automatically updates existing chromecast entities.
+    """
+    if chromecast.uuid is None:
+        # Found a cast without UUID, we don't store it because we won't be able
+        # to update it anyway.
+        return CastDevice(chromecast)
+
+    # Found a cast with UUID
+    added_casts = hass.data[ADDED_CAST_DEVICES_KEY]
+    old_cast_device = added_casts.get(chromecast.uuid)
+    if old_cast_device is None:
+        # -> New cast device
+        cast_device = CastDevice(chromecast)
+        added_casts[chromecast.uuid] = cast_device
+        return cast_device
+
+    old_key = (old_cast_device.cast.host,
+               old_cast_device.cast.port,
+               old_cast_device.cast.uuid)
+    new_key = (chromecast.host, chromecast.port, chromecast.uuid)
+
+    if old_key == new_key:
+        # Re-discovered with same data, ignore
+        return None
+
+    # -> Cast device changed host
+    # Remove old pychromecast.Chromecast from global list, because it isn't
+    # valid anymore
+    old_cast_device.async_set_chromecast(chromecast)
+    return None
+
+
+@asyncio.coroutine
+def async_setup_platform(hass: HomeAssistantType, config: ConfigType,
+                         async_add_devices, discovery_info=None):
     """Set up the cast platform."""
     import pychromecast
 
     # Import CEC IGNORE attributes
     pychromecast.IGNORE_CEC += config.get(CONF_IGNORE_CEC, [])
+    hass.data.setdefault(ADDED_CAST_DEVICES_KEY, {})
+    hass.data.setdefault(KNOWN_CHROMECASTS_KEY, {})
 
-    known_hosts = hass.data.get(KNOWN_HOSTS_KEY)
+    # None -> use discovery; (host, port) -> manually specify chromecast.
-    if known_hosts is None:
+    want_host = None
-        known_hosts = hass.data[KNOWN_HOSTS_KEY] = []
-
     if discovery_info:
-        host = (discovery_info.get('host'), discovery_info.get('port'))
+        want_host = (discovery_info.get('host'), discovery_info.get('port'))
-
-        if host in known_hosts:
-            return
-
-        hosts = [host]
-
     elif CONF_HOST in config:
-        host = (config.get(CONF_HOST), DEFAULT_PORT)
+        want_host = (config.get(CONF_HOST), DEFAULT_PORT)
 
-        if host in known_hosts:
+    enable_discovery = False
-            return
+    if want_host is None:
+        # We were explicitly told to enable pychromecast discovery.
+        enable_discovery = True
+    elif want_host[1] != DEFAULT_PORT:
+        # We're trying to add a group, so we have to use pychromecast's
+        # discovery to get the correct friendly name.
+        enable_discovery = True
 
-        hosts = [host]
+    if enable_discovery:
+        @callback
+        def async_cast_discovered(chromecast):
+            """Callback for when a new chromecast is discovered."""
+            if want_host is not None and \
+                    (chromecast.host, chromecast.port) != want_host:
+                return  # for groups, only add requested device
+            cast_device = _async_create_cast_device(hass, chromecast)
 
+            if cast_device is not None:
+                async_add_devices([cast_device])
+
+        async_dispatcher_connect(hass, SIGNAL_CAST_DISCOVERED,
+                                 async_cast_discovered)
+        # Re-play the callback for all past chromecasts, store the objects in
+        # a list to avoid concurrent modification resulting in exception.
+        for chromecast in list(hass.data[KNOWN_CHROMECASTS_KEY].values()):
+            async_cast_discovered(chromecast)
+
+        hass.async_add_job(_setup_internal_discovery, hass)
     else:
-        hosts = [tuple(dev[:2]) for dev in pychromecast.discover_chromecasts()
+        # Manually add a "normal" Chromecast, we can do that without discovery.
-                 if tuple(dev[:2]) not in known_hosts]
+        try:
-
+            chromecast = pychromecast.Chromecast(*want_host)
-    casts = []
+        except pychromecast.ChromecastConnectionError:
-
+            _LOGGER.warning("Can't set up chromecast on %s", want_host[0])
-    # get_chromecasts() returns Chromecast objects with the correct friendly
+            raise
-    # name for grouped devices
+        key = (chromecast.host, chromecast.port, chromecast.uuid)
-    all_chromecasts = pychromecast.get_chromecasts()
+        cast_device = _async_create_cast_device(hass, chromecast)
-
+        if cast_device is not None:
-    for host in hosts:
+            hass.data[KNOWN_CHROMECASTS_KEY][key] = chromecast
-        (_, port) = host
+            async_add_devices([cast_device])
-        found = [device for device in all_chromecasts
-                 if (device.host, device.port) == host]
-        if found:
-            try:
-                casts.append(CastDevice(found[0]))
-                known_hosts.append(host)
-            except pychromecast.ChromecastConnectionError:
-                pass
-
-        # do not add groups using pychromecast.Chromecast as it leads to names
-        # collision since pychromecast.Chromecast will get device name instead
-        # of group name
-        elif port == DEFAULT_PORT:
-            try:
-                # add the device anyway, get_chromecasts couldn't find it
-                casts.append(CastDevice(pychromecast.Chromecast(*host)))
-                known_hosts.append(host)
-            except pychromecast.ChromecastConnectionError:
-                pass
-
-    add_devices(casts)
 
 
 class CastDevice(MediaPlayerDevice):
@@ -109,16 +198,13 @@ class CastDevice(MediaPlayerDevice):
 
     def __init__(self, chromecast):
         """Initialize the Cast device."""
-        self.cast = chromecast
+        self.cast = None  # type: pychromecast.Chromecast
-
+        self.cast_status = None
-        self.cast.socket_client.receiver_controller.register_status_listener(
+        self.media_status = None
-            self)
-        self.cast.socket_client.media_controller.register_status_listener(self)
-
-        self.cast_status = self.cast.status
-        self.media_status = self.cast.media_controller.status
         self.media_status_received = None
 
+        self.async_set_chromecast(chromecast)
+
     @property
     def should_poll(self):
         """No polling needed."""
@@ -325,3 +411,39 @@ class CastDevice(MediaPlayerDevice):
         self.media_status = status
         self.media_status_received = dt_util.utcnow()
         self.schedule_update_ha_state()
+
+    @property
+    def unique_id(self) -> str:
+        """Return an unique ID."""
+        if self.cast.uuid is not None:
+            return str(self.cast.uuid)
+        return None
+
+    @callback
+    def async_set_chromecast(self, chromecast):
+        """Set the internal Chromecast object and disconnect the previous."""
+        self._async_disconnect()
+
+        self.cast = chromecast
+
+        self.cast.socket_client.receiver_controller.register_status_listener(
+            self)
+        self.cast.socket_client.media_controller.register_status_listener(self)
+
+        self.cast_status = self.cast.status
+        self.media_status = self.cast.media_controller.status
+
+    @asyncio.coroutine
+    def async_will_remove_from_hass(self):
+        """Disconnect Chromecast object when removed."""
+        self._async_disconnect()
+
+    @callback
+    def _async_disconnect(self):
+        """Disconnect Chromecast object if it is set."""
+        if self.cast is None:
+            return
+        _LOGGER.debug("Disconnecting existing chromecast object")
+        old_key = (self.cast.host, self.cast.port, self.cast.uuid)
+        self.hass.data[KNOWN_CHROMECASTS_KEY].pop(old_key)
+        self.cast.disconnect(blocking=False)

tests/components/media_player/test_cast.py

@@ -1,12 +1,15 @@
 """The tests for the Cast Media player platform."""
 # pylint: disable=protected-access
-import unittest
+import asyncio
-from unittest.mock import patch, MagicMock
+from typing import Optional
+from unittest.mock import patch, MagicMock, Mock
+from uuid import UUID
 
 import pytest
 
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.components.media_player import cast
-from tests.common import get_test_home_assistant
 
 
 @pytest.fixture(autouse=True)
@@ -18,83 +21,221 @@ def cast_mock():
         yield
 
 
-class FakeChromeCast(object):
+# pylint: disable=invalid-name
-    """A fake Chrome Cast."""
+FakeUUID = UUID('57355bce-9364-4aa6-ac1e-eb849dccf9e2')
-
-    def __init__(self, host, port):
-        """Initialize the fake Chrome Cast."""
-        self.host = host
-        self.port = port
 
 
-class TestCastMediaPlayer(unittest.TestCase):
+def get_fake_chromecast(host='192.168.178.42', port=8009,
-    """Test the media_player module."""
+                        uuid: Optional[UUID] = FakeUUID):
+    """Generate a Fake Chromecast object with the specified arguments."""
+    return MagicMock(host=host, port=port, uuid=uuid)
 
-    def setUp(self):
-        """Setup things to be run when tests are started."""
-        self.hass = get_test_home_assistant()
 
-    def tearDown(self):
+@asyncio.coroutine
-        """Stop everything that was started."""
+def async_setup_cast(hass, config=None, discovery_info=None):
-        self.hass.stop()
+    """Helper to setup the cast platform."""
+    if config is None:
+        config = {}
+    add_devices = Mock()
 
-    @patch('homeassistant.components.media_player.cast.CastDevice')
+    yield from cast.async_setup_platform(hass, config, add_devices,
-    @patch('pychromecast.get_chromecasts')
+                                         discovery_info=discovery_info)
-    def test_filter_duplicates(self, mock_get_chromecasts, mock_device):
+    yield from hass.async_block_till_done()
-        """Test filtering of duplicates."""
-        mock_get_chromecasts.return_value = [
-            FakeChromeCast('some_host', cast.DEFAULT_PORT)
-        ]
 
-        # Test chromecasts as if they were hardcoded in configuration.yaml
+    return add_devices
-        cast.setup_platform(self.hass, {
-            'host': 'some_host'
-        }, lambda _: _)
 
-        assert mock_device.called
 
-        mock_device.reset_mock()
+@asyncio.coroutine
-        assert not mock_device.called
+def async_setup_cast_internal_discovery(hass, config=None,
+                                        discovery_info=None,
+                                        no_from_host_patch=False):
+    """Setup the cast platform and the discovery."""
+    listener = MagicMock(services={})
 
-        # Test chromecasts as if they were automatically discovered
+    with patch('pychromecast.start_discovery',
-        cast.setup_platform(self.hass, {}, lambda _: _, {
+               return_value=(listener, None)) as start_discovery:
-            'host': 'some_host',
+        add_devices = yield from async_setup_cast(hass, config, discovery_info)
-            'port': cast.DEFAULT_PORT,
+        yield from hass.async_block_till_done()
-        })
+        yield from hass.async_block_till_done()
-        assert not mock_device.called
 
-    @patch('homeassistant.components.media_player.cast.CastDevice')
+        assert start_discovery.call_count == 1
-    @patch('pychromecast.get_chromecasts')
-    @patch('pychromecast.Chromecast')
-    def test_fallback_cast(self, mock_chromecast, mock_get_chromecasts,
-                           mock_device):
-        """Test falling back to creating Chromecast when not discovered."""
-        mock_get_chromecasts.return_value = [
-            FakeChromeCast('some_host', cast.DEFAULT_PORT)
-        ]
 
-        # Test chromecasts as if they were hardcoded in configuration.yaml
+        discovery_callback = start_discovery.call_args[0][0]
-        cast.setup_platform(self.hass, {
-            'host': 'some_other_host'
-        }, lambda _: _)
 
-        assert mock_chromecast.called
+    def discover_chromecast(service_name, chromecast):
-        assert mock_device.called
+        """Discover a chromecast device."""
+        listener.services[service_name] = (
+            chromecast.host, chromecast.port, chromecast.uuid, None, None)
+        if no_from_host_patch:
+            discovery_callback(service_name)
+        else:
+            with patch('pychromecast._get_chromecast_from_host',
+                       return_value=chromecast):
+                discovery_callback(service_name)
 
-    @patch('homeassistant.components.media_player.cast.CastDevice')
+    return discover_chromecast, add_devices
-    @patch('pychromecast.get_chromecasts')
-    @patch('pychromecast.Chromecast')
-    def test_fallback_cast_group(self, mock_chromecast, mock_get_chromecasts,
-                                 mock_device):
-        """Test not creating Cast Group when not discovered."""
-        mock_get_chromecasts.return_value = [
-            FakeChromeCast('some_host', cast.DEFAULT_PORT)
-        ]
 
-        # Test chromecasts as if they were automatically discovered
+
-        cast.setup_platform(self.hass, {}, lambda _: _, {
+@asyncio.coroutine
-            'host': 'some_other_host',
+def test_start_discovery_called_once(hass):
-            'port': 43546,
+    """Test pychromecast.start_discovery called exactly once."""
-        })
+    with patch('pychromecast.start_discovery',
-        assert not mock_chromecast.called
+               return_value=(None, None)) as start_discovery:
-        assert not mock_device.called
+        yield from async_setup_cast(hass)
+
+        assert start_discovery.call_count == 1
+
+        yield from async_setup_cast(hass)
+        assert start_discovery.call_count == 1
+
+
+@asyncio.coroutine
+def test_stop_discovery_called_on_stop(hass):
+    """Test pychromecast.stop_discovery called on shutdown."""
+    with patch('pychromecast.start_discovery',
+               return_value=(None, 'the-browser')) as start_discovery:
+        yield from async_setup_cast(hass)
+
+        assert start_discovery.call_count == 1
+
+    with patch('pychromecast.stop_discovery') as stop_discovery:
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
+        yield from hass.async_block_till_done()
+
+        stop_discovery.assert_called_once_with('the-browser')
+
+    with patch('pychromecast.start_discovery',
+               return_value=(None, 'the-browser')) as start_discovery:
+        yield from async_setup_cast(hass)
+
+        assert start_discovery.call_count == 1
+
+
+@asyncio.coroutine
+def test_internal_discovery_callback_only_generates_once(hass):
+    """Test _get_chromecast_from_host only called once per device."""
+    discover_cast, _ = yield from async_setup_cast_internal_discovery(
+        hass, no_from_host_patch=True)
+    chromecast = get_fake_chromecast()
+
+    with patch('pychromecast._get_chromecast_from_host',
+               return_value=chromecast) as gen_chromecast:
+        discover_cast('the-service', chromecast)
+        mdns = (chromecast.host, chromecast.port, chromecast.uuid, None, None)
+        gen_chromecast.assert_called_once_with(mdns, blocking=True)
+
+        discover_cast('the-service', chromecast)
+        gen_chromecast.reset_mock()
+        assert gen_chromecast.call_count == 0
+
+
+@asyncio.coroutine
+def test_internal_discovery_callback_calls_dispatcher(hass):
+    """Test internal discovery calls dispatcher."""
+    discover_cast, _ = yield from async_setup_cast_internal_discovery(hass)
+    chromecast = get_fake_chromecast()
+
+    with patch('pychromecast._get_chromecast_from_host',
+               return_value=chromecast):
+        signal = MagicMock()
+
+        async_dispatcher_connect(hass, 'cast_discovered', signal)
+        discover_cast('the-service', chromecast)
+        yield from hass.async_block_till_done()
+
+        signal.assert_called_once_with(chromecast)
+
+
+@asyncio.coroutine
+def test_internal_discovery_callback_with_connection_error(hass):
+    """Test internal discovery not calling dispatcher on ConnectionError."""
+    import pychromecast  # imports mock pychromecast
+
+    pychromecast.ChromecastConnectionError = IOError
+
+    discover_cast, _ = yield from async_setup_cast_internal_discovery(
+        hass, no_from_host_patch=True)
+    chromecast = get_fake_chromecast()
+
+    with patch('pychromecast._get_chromecast_from_host',
+               side_effect=pychromecast.ChromecastConnectionError):
+        signal = MagicMock()
+
+        async_dispatcher_connect(hass, 'cast_discovered', signal)
+        discover_cast('the-service', chromecast)
+        yield from hass.async_block_till_done()
+
+        assert signal.call_count == 0
+
+
+def test_create_cast_device_without_uuid(hass):
+    """Test create a cast device without a UUID."""
+    chromecast = get_fake_chromecast(uuid=None)
+    cast_device = cast._async_create_cast_device(hass, chromecast)
+    assert cast_device is not None
+
+
+def test_create_cast_device_with_uuid(hass):
+    """Test create cast devices with UUID."""
+    added_casts = hass.data[cast.ADDED_CAST_DEVICES_KEY] = {}
+    chromecast = get_fake_chromecast()
+    cast_device = cast._async_create_cast_device(hass, chromecast)
+    assert cast_device is not None
+    assert chromecast.uuid in added_casts
+
+    with patch.object(cast_device, 'async_set_chromecast') as mock_set:
+        assert cast._async_create_cast_device(hass, chromecast) is None
+        assert mock_set.call_count == 0
+
+        chromecast = get_fake_chromecast(host='192.168.178.1')
+        assert cast._async_create_cast_device(hass, chromecast) is None
+        assert mock_set.call_count == 1
+        mock_set.assert_called_once_with(chromecast)
+
+
+@asyncio.coroutine
+def test_normal_chromecast_not_starting_discovery(hass):
+    """Test cast platform not starting discovery when not required."""
+    chromecast = get_fake_chromecast()
+
+    with patch('pychromecast.Chromecast', return_value=chromecast):
+        add_devices = yield from async_setup_cast(hass, {'host': 'host1'})
+        assert add_devices.call_count == 1
+
+        # Same entity twice
+        add_devices = yield from async_setup_cast(hass, {'host': 'host1'})
+        assert add_devices.call_count == 0
+
+        hass.data[cast.ADDED_CAST_DEVICES_KEY] = {}
+        add_devices = yield from async_setup_cast(
+            hass, discovery_info={'host': 'host1', 'port': 8009})
+        assert add_devices.call_count == 1
+
+        hass.data[cast.ADDED_CAST_DEVICES_KEY] = {}
+        add_devices = yield from async_setup_cast(
+            hass, discovery_info={'host': 'host1', 'port': 42})
+        assert add_devices.call_count == 0
+
+
+@asyncio.coroutine
+def test_replay_past_chromecasts(hass):
+    """Test cast platform re-playing past chromecasts when adding new one."""
+    cast_group1 = get_fake_chromecast(host='host1', port=42)
+    cast_group2 = get_fake_chromecast(host='host2', port=42, uuid=UUID(
+        '9462202c-e747-4af5-a66b-7dce0e1ebc09'))
+
+    discover_cast, add_dev1 = yield from async_setup_cast_internal_discovery(
+        hass, discovery_info={'host': 'host1', 'port': 42})
+    discover_cast('service2', cast_group2)
+    yield from hass.async_block_till_done()
+    assert add_dev1.call_count == 0
+
+    discover_cast('service1', cast_group1)
+    yield from hass.async_block_till_done()
+    yield from hass.async_block_till_done()  # having jobs that add jobs
+    assert add_dev1.call_count == 1
+
+    add_dev2 = yield from async_setup_cast(
+        hass, discovery_info={'host': 'host2', 'port': 42})
+    yield from hass.async_block_till_done()
+    assert add_dev2.call_count == 1