From e14893416fb98d3ef9a37d816e0ee719105b33a1 Mon Sep 17 00:00:00 2001
From: Otto Winter <otto@otto-winter.com>
Date: Sat, 24 Feb 2018 19:27:44 +0100
Subject: [PATCH] Cast automatically drop connection (#12635)

---
 homeassistant/components/media_player/cast.py | 17 +++++++++++------
 tests/components/media_player/test_cast.py    | 12 +++++++++++-
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/media_player/cast.py b/homeassistant/components/media_player/cast.py
index 40e09ea328c..d3cf2f7b501 100644
--- a/homeassistant/components/media_player/cast.py
+++ b/homeassistant/components/media_player/cast.py
@@ -8,9 +8,11 @@ https://home-assistant.io/components/media_player.cast/
 import asyncio
 import logging
 import threading
+import functools
 
 import voluptuous as vol
 
+from homeassistant.exceptions import PlatformNotReady
 from homeassistant.helpers.typing import HomeAssistantType, ConfigType
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import (dispatcher_send,
@@ -34,6 +36,7 @@ CONF_IGNORE_CEC = 'ignore_cec'
 CAST_SPLASH = 'https://home-assistant.io/images/cast/splash.png'
 
 DEFAULT_PORT = 8009
+SOCKET_CLIENT_RETRIES = 10
 
 SUPPORT_CAST = SUPPORT_PAUSE | SUPPORT_VOLUME_SET | SUPPORT_VOLUME_MUTE | \
     SUPPORT_TURN_ON | SUPPORT_TURN_OFF | SUPPORT_PREVIOUS_TRACK | \
@@ -76,7 +79,7 @@ def _setup_internal_discovery(hass: HomeAssistantType) -> None:
         try:
             # pylint: disable=protected-access
             chromecast = pychromecast._get_chromecast_from_host(
-                mdns, blocking=True)
+                mdns, blocking=True, tries=SOCKET_CLIENT_RETRIES)
         except pychromecast.ChromecastConnectionError:
             _LOGGER.debug("Can't set up cast with mDNS info %s. "
                           "Assuming it's not a Chromecast", mdns)
@@ -182,11 +185,13 @@ def async_setup_platform(hass: HomeAssistantType, config: ConfigType,
     else:
         # Manually add a "normal" Chromecast, we can do that without discovery.
         try:
-            chromecast = yield from hass.async_add_job(
-                pychromecast.Chromecast, *want_host)
-        except pychromecast.ChromecastConnectionError:
-            _LOGGER.warning("Can't set up chromecast on %s", want_host[0])
-            raise
+            func = functools.partial(pychromecast.Chromecast, *want_host,
+                                     tries=SOCKET_CLIENT_RETRIES)
+            chromecast = yield from hass.async_add_job(func)
+        except pychromecast.ChromecastConnectionError as err:
+            _LOGGER.warning("Can't set up chromecast on %s: %s",
+                            want_host[0], err)
+            raise PlatformNotReady
         key = (chromecast.host, chromecast.port, chromecast.uuid)
         cast_device = _async_create_cast_device(hass, chromecast)
         if cast_device is not None:
diff --git a/tests/components/media_player/test_cast.py b/tests/components/media_player/test_cast.py
index 6eeb9136b07..aaaad47d8dc 100644
--- a/tests/components/media_player/test_cast.py
+++ b/tests/components/media_player/test_cast.py
@@ -7,6 +7,7 @@ from uuid import UUID
 
 import pytest
 
+from homeassistant.exceptions import PlatformNotReady
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.components.media_player import cast
@@ -122,7 +123,7 @@ def test_internal_discovery_callback_only_generates_once(hass):
                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)
+        gen_chromecast.assert_called_once_with(mdns, blocking=True, tries=10)
 
         discover_cast('the-service', chromecast)
         gen_chromecast.reset_mock()
@@ -196,6 +197,10 @@ def test_create_cast_device_with_uuid(hass):
 @asyncio.coroutine
 def test_normal_chromecast_not_starting_discovery(hass):
     """Test cast platform not starting discovery when not required."""
+    import pychromecast  # imports mock pychromecast
+
+    pychromecast.ChromecastConnectionError = IOError
+
     chromecast = get_fake_chromecast()
 
     with patch('pychromecast.Chromecast', return_value=chromecast):
@@ -216,6 +221,11 @@ def test_normal_chromecast_not_starting_discovery(hass):
             hass, discovery_info={'host': 'host1', 'port': 42})
         assert add_devices.call_count == 0
 
+    with patch('pychromecast.Chromecast',
+               side_effect=pychromecast.ChromecastConnectionError):
+        with pytest.raises(PlatformNotReady):
+            yield from async_setup_cast(hass, {'host': 'host3'})
+
 
 @asyncio.coroutine
 def test_replay_past_chromecasts(hass):