Finish all tasks before setup phase is done (#4606)

homeassistant/bootstrap.py

@@ -365,6 +365,7 @@ def async_from_config_dict(config: Dict[str, Any],
     Dynamically loads required components and its dependencies.
     This method is a coroutine.
     """
+    hass.async_track_tasks()
     setup_lock = hass.data.get('setup_lock')
     if setup_lock is None:
         setup_lock = hass.data['setup_lock'] = asyncio.Lock(loop=hass.loop)
@@ -427,6 +428,8 @@ def async_from_config_dict(config: Dict[str, Any],
 
     setup_lock.release()
 
+    yield from hass.async_stop_track_tasks()
+
     return hass
 
 

homeassistant/core.py

@@ -236,6 +236,12 @@ class HomeAssistant(object):
         """Track tasks so you can wait for all tasks to be done."""
         self.async_add_job = self._async_add_job_tracking
 
+    @asyncio.coroutine
+    def async_stop_track_tasks(self):
+        """Track tasks so you can wait for all tasks to be done."""
+        yield from self.async_block_till_done()
+        self.async_add_job = self._async_add_job
+
     @callback
     def async_run_job(self, target: Callable[..., None], *args: Any) -> None:
         """Run a job from within the event loop.

tests/helpers/test_discovery.py

@@ -166,7 +166,8 @@ class TestHelpersDiscovery:
 
         def component1_setup(hass, config):
             """Setup mock component."""
-            discovery.discover(hass, 'test_component2')
+            discovery.discover(hass, 'test_component2',
+                               component='test_component2')
             return True
 
         def component2_setup(hass, config):

tests/test_bootstrap.py

@@ -1,20 +1,26 @@
 """Test the bootstrapping."""
 # pylint: disable=protected-access
+import os
 from unittest import mock
 import threading
 import logging
 
 import voluptuous as vol
 
+from homeassistant.core import callback
+from homeassistant.const import EVENT_HOMEASSISTANT_START
+import homeassistant.config as config_util
 from homeassistant import bootstrap, loader
 import homeassistant.util.dt as dt_util
 from homeassistant.helpers.config_validation import PLATFORM_SCHEMA
+from homeassistant.helpers import discovery
 
 from tests.common import \
     get_test_home_assistant, MockModule, MockPlatform, \
-    assert_setup_component, patch_yaml_files
+    assert_setup_component, patch_yaml_files, get_test_config_dir
 
 ORIG_TIMEZONE = dt_util.DEFAULT_TIME_ZONE
+VERSION_PATH = os.path.join(get_test_config_dir(), config_util.VERSION_FILE)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -43,6 +49,8 @@ class TestBootstrap:
         dt_util.DEFAULT_TIME_ZONE = ORIG_TIMEZONE
         self.hass.stop()
         loader._COMPONENT_CACHE = self.backup_cache
+        if os.path.isfile(VERSION_PATH):
+            os.remove(VERSION_PATH)
 
     @mock.patch(
         # prevent .HA_VERISON file from being written
@@ -381,3 +389,48 @@ class TestBootstrap:
         assert bootstrap.setup_component(self.hass, 'disabled_component')
         assert loader.get_component('disabled_component') is not None
         assert 'disabled_component' in self.hass.config.components
+
+    def test_all_work_done_before_start(self):
+        """Test all init work done till start."""
+        call_order = []
+
+        def component1_setup(hass, config):
+            """Setup mock component."""
+            discovery.discover(hass, 'test_component2',
+                               component='test_component2')
+            discovery.discover(hass, 'test_component3',
+                               component='test_component3')
+            return True
+
+        def component_track_setup(hass, config):
+            """Setup mock component."""
+            call_order.append(1)
+            return True
+
+        loader.set_component(
+            'test_component1',
+            MockModule('test_component1', setup=component1_setup))
+
+        loader.set_component(
+            'test_component2',
+            MockModule('test_component2', setup=component_track_setup))
+
+        loader.set_component(
+            'test_component3',
+            MockModule('test_component3', setup=component_track_setup))
+
+        @callback
+        def track_start(event):
+            """Track start event."""
+            call_order.append(2)
+
+        self.hass.bus.listen_once(EVENT_HOMEASSISTANT_START, track_start)
+
+        self.hass.loop.run_until_complete = \
+            lambda _: self.hass.block_till_done()
+
+        bootstrap.from_config_dict({'test_component1': None}, self.hass)
+
+        self.hass.start()
+
+        assert call_order == [1, 1, 2]