From 7055fddfb47cc2a10fb1d4e8cb7da66cf0beb069 Mon Sep 17 00:00:00 2001
From: Anton Sarukhanov <code@ant.sr>
Date: Tue, 23 May 2017 14:29:27 -0700
Subject: [PATCH] Don't block startup more than 60 seconds while waiting for
 components. (#7739)

---
 homeassistant/helpers/entity_component.py | 15 +++++++++----
 tests/common.py                           |  5 ++++-
 tests/helpers/test_entity_component.py    | 26 ++++++++++++++++++++++-
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/homeassistant/helpers/entity_component.py b/homeassistant/helpers/entity_component.py
index 2c1801a63427..9ebad3862f3d 100644
--- a/homeassistant/helpers/entity_component.py
+++ b/homeassistant/helpers/entity_component.py
@@ -19,6 +19,7 @@ from homeassistant.util.async import (
 
 DEFAULT_SCAN_INTERVAL = timedelta(seconds=15)
 SLOW_SETUP_WARNING = 10
+SLOW_SETUP_MAX_WAIT = 60
 
 
 class EntityComponent(object):
@@ -145,20 +146,26 @@ class EntityComponent(object):
 
         try:
             if getattr(platform, 'async_setup_platform', None):
-                yield from platform.async_setup_platform(
+                task = platform.async_setup_platform(
                     self.hass, platform_config,
                     entity_platform.async_schedule_add_entities, discovery_info
                 )
             else:
-                yield from self.hass.loop.run_in_executor(
+                task = self.hass.loop.run_in_executor(
                     None, platform.setup_platform, self.hass, platform_config,
                     entity_platform.schedule_add_entities, discovery_info
                 )
-
+            yield from asyncio.wait_for(
+                asyncio.shield(task, loop=self.hass.loop),
+                SLOW_SETUP_MAX_WAIT, loop=self.hass.loop)
             yield from entity_platform.async_block_entities_done()
-
             self.hass.config.components.add(
                 '{}.{}'.format(self.domain, platform_type))
+        except asyncio.TimeoutError:
+            self.logger.error(
+                "Setup of platform %s is taking longer than %s seconds."
+                " Startup will proceed without waiting any longer.",
+                platform_type, SLOW_SETUP_MAX_WAIT)
         except Exception:  # pylint: disable=broad-except
             self.logger.exception(
                 "Error while setting up platform %s", platform_type)
diff --git a/tests/common.py b/tests/common.py
index 30bd772a81f0..735b1dfce98d 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -317,7 +317,7 @@ class MockPlatform(object):
 
     # pylint: disable=invalid-name
     def __init__(self, setup_platform=None, dependencies=None,
-                 platform_schema=None):
+                 platform_schema=None, async_setup_platform=None):
         """Initialize the platform."""
         self.DEPENDENCIES = dependencies or []
         self._setup_platform = setup_platform
@@ -325,6 +325,9 @@ class MockPlatform(object):
         if platform_schema is not None:
             self.PLATFORM_SCHEMA = platform_schema
 
+        if async_setup_platform is not None:
+            self.async_setup_platform = async_setup_platform
+
     def setup_platform(self, hass, config, add_devices, discovery_info=None):
         """Set up the platform."""
         if self._setup_platform is not None:
diff --git a/tests/helpers/test_entity_component.py b/tests/helpers/test_entity_component.py
index ade8c4ebd8ad..566306d7fe73 100644
--- a/tests/helpers/test_entity_component.py
+++ b/tests/helpers/test_entity_component.py
@@ -13,6 +13,7 @@ from homeassistant.components import group
 from homeassistant.helpers.entity import Entity, generate_entity_id
 from homeassistant.helpers.entity_component import (
     EntityComponent, DEFAULT_SCAN_INTERVAL, SLOW_SETUP_WARNING)
+from homeassistant.helpers import entity_component
 
 from homeassistant.helpers import discovery
 import homeassistant.util.dt as dt_util
@@ -472,7 +473,6 @@ def test_platform_warn_slow_setup(hass):
             }
         })
         assert mock_call.called
-        assert len(mock_call.mock_calls) == 2
 
         timeout, logger_method = mock_call.mock_calls[0][1][:2]
 
@@ -482,6 +482,30 @@ def test_platform_warn_slow_setup(hass):
         assert mock_call().cancel.called
 
 
+@asyncio.coroutine
+def test_platform_error_slow_setup(hass, caplog):
+    """Don't block startup more than SLOW_SETUP_MAX_WAIT."""
+    with patch.object(entity_component, 'SLOW_SETUP_MAX_WAIT', 0):
+        called = []
+
+        @asyncio.coroutine
+        def setup_platform(*args):
+            called.append(1)
+            yield from asyncio.sleep(1, loop=hass.loop)
+
+        platform = MockPlatform(async_setup_platform=setup_platform)
+        component = EntityComponent(_LOGGER, DOMAIN, hass)
+        loader.set_component('test_domain.test_platform', platform)
+        yield from component.async_setup({
+            DOMAIN: {
+                'platform': 'test_platform',
+            }
+        })
+        assert len(called) == 1
+        assert 'test_domain.test_platform' not in hass.config.components
+        assert 'test_platform is taking longer than 0 seconds' in caplog.text
+
+
 @asyncio.coroutine
 def test_extract_from_service_available_device(hass):
     """Test the extraction of entity from service and device is available."""