home-assistant · balloob · Oct 22, 2017 · Oct 22, 2017 · Oct 22, 2017 · Oct 22, 2017
@@ -200,34 +200,11 @@ def async_update_ha_state(self, force_refresh=False):
 
         # update entity data
         if force_refresh:
-            if self._update_staged:
-                return
-            self._update_staged = True
-
-            # Process update sequential
-            if self.parallel_updates:
-                yield from self.parallel_updates.acquire()
-
-            update_warn = self.hass.loop.call_later(
-                SLOW_UPDATE_WARNING, _LOGGER.warning,
-                "Update of %s is taking over %s seconds", self.entity_id,
-                SLOW_UPDATE_WARNING
-            )
-
             try:
-                if hasattr(self, 'async_update'):
-                    # pylint: disable=no-member
-                    yield from self.async_update()
-                else:
-                    yield from self.hass.async_add_job(self.update)
+                yield from self.async_device_update()
             except Exception:  # pylint: disable=broad-except
                 _LOGGER.exception("Update for %s fails", self.entity_id)
                 return
-            finally:
-                self._update_staged = False
-                update_warn.cancel()
-                if self.parallel_updates:
-                    self.parallel_updates.release()
 
         start = timer()
 
@@ -304,6 +281,39 @@ def async_schedule_update_ha_state(self, force_refresh=False):
         """Schedule a update ha state change task."""
         self.hass.async_add_job(self.async_update_ha_state(force_refresh))
 
+    def async_device_update(self, warning=True):
+        """Process 'update' or 'async_update' from entity.
+
+        This method is a coroutine.
+        """
+        if self._update_staged:
+            return
+        self._update_staged = True
+
+        # Process update sequential
+        if self.parallel_updates:
+            yield from self.parallel_updates.acquire()
+
+        if warning:
+            update_warn = self.hass.loop.call_later(
+                SLOW_UPDATE_WARNING, _LOGGER.warning,
+                "Update of %s is taking over %s seconds", self.entity_id,
+                SLOW_UPDATE_WARNING
+            )
+
+        try:
+            if hasattr(self, 'async_update'):
+                # pylint: disable=no-member
+                yield from self.async_update()
+            else:
+                yield from self.hass.async_add_job(self.update)
+        finally:
+            self._update_staged = False
+            if warning:
+                update_warn.cancel()
+            if self.parallel_updates:
+                self.parallel_updates.release()
+
     def remove(self) -> None:
         """Remove entity from HASS."""
         run_coroutine_threadsafe(

@@ -210,6 +210,15 @@ def async_add_entity(self, entity, platform=None, update_before_add=False):
 
         entity.hass = self.hass
 
+        # Update properties before we generate the entity_id
+        if update_before_add:
+            try:
+                yield from entity.async_device_update(warning=False)
+            except Exception:  # pylint: disable=broad-except
+                self.logger.exception("Error on device update!")
+                return False
+
+        # Write entity_id to entity
         if getattr(entity, 'entity_id', None) is None:
             object_id = entity.name or DEVICE_DEFAULT_NAME
 
@@ -234,7 +243,7 @@ def async_add_entity(self, entity, platform=None, update_before_add=False):
         if hasattr(entity, 'async_added_to_hass'):
             yield from entity.async_added_to_hass()
 
-        yield from entity.async_update_ha_state(update_before_add)
+        yield from entity.async_update_ha_state()
 
         return True
 
@@ -361,12 +370,14 @@ def async_schedule_add_entities(self, new_entities,
 
     def add_entities(self, new_entities, update_before_add=False):
         """Add entities for a single platform."""
+        # That avoid deadlocks
         if update_before_add:
-            for entity in new_entities:
-                entity.update()
+            self.component.logger.warning(
+                "Call 'add_entities' with update_before_add=True "
+                "only inside tests or you can run into a deadlock!")
 
         run_coroutine_threadsafe(
-            self.async_add_entities(list(new_entities), False),
+            self.async_add_entities(list(new_entities), update_before_add),
             self.component.hass.loop).result()
 
     @asyncio.coroutine

diff --git a/tests/helpers/test_entity.py b/tests/helpers/test_entity.py
@@ -193,6 +193,30 @@ def async_update():
         assert update_call
 
 
+@asyncio.coroutine
+def test_warn_slow_device_update_disabled(hass):
+    """Disable slow update warning with async_device_update."""
+    update_call = False
+
+    @asyncio.coroutine
+    def async_update():
+        """Mock async update."""
+        nonlocal update_call
+        update_call = True
+
+    mock_entity = entity.Entity()
+    mock_entity.hass = hass
+    mock_entity.entity_id = 'comp_test.test_entity'
+    mock_entity.async_update = async_update
+
+    with patch.object(hass.loop, 'call_later', MagicMock()) \
+            as mock_call:
+        yield from mock_entity.async_device_update(warning=False)
+
+        assert not mock_call.called
+        assert update_call
+
+
 @asyncio.coroutine
 def test_async_schedule_update_ha_state(hass):
     """Warn we log when entity update takes a long time and trow exception."""

diff --git a/tests/helpers/test_entity_component.py b/tests/helpers/test_entity_component.py
@@ -208,30 +208,6 @@ def test_update_state_adds_entities_with_update_befor_add_false(self):
         assert 1 == len(self.hass.states.entity_ids())
         assert not ent.update.called
 
-    def test_adds_entities_with_update_befor_add_true_deadlock_protect(self):
-        """Test if call update before add to state machine.
-
-        It need to run update inside executor and never call
-        async_add_entities with True
-        """
-        call = []
-        component = EntityComponent(_LOGGER, DOMAIN, self.hass)
-
-        @asyncio.coroutine
-        def async_add_entities_fake(entities, update_befor_add):
-            """Fake add_entities_call."""
-            call.append(update_befor_add)
-        component._platforms['core'].async_add_entities = \
-            async_add_entities_fake
-
-        ent = EntityTest()
-        ent.update = Mock(spec_set=True)
-        component.add_entities([ent], True)
-
-        assert ent.update.called
-        assert len(call) == 1
-        assert not call[0]
-
     def test_not_adding_duplicate_entities(self):
         """Test for not adding duplicate entities."""
         component = EntityComponent(_LOGGER, DOMAIN, self.hass)
@@ -654,3 +630,24 @@ def test_pararell_updates_sync_platform(hass):
     handle = list(component._platforms.values())[-1]
 
     assert handle.parallel_updates is not None
+
+
+@asyncio.coroutine
+def test_raise_error_on_update(hass):
+    """Test the add entity if they raise an error on update."""
+    updates = []
+    component = EntityComponent(_LOGGER, DOMAIN, hass)
+    entity1 = EntityTest(name='test_1')
+    entity2 = EntityTest(name='test_2')
+
+    def _raise():
+        """Helper to raise a exception."""
+        raise AssertionError
+    
+    entity1.update = _raise
+    entity2.update = lambda: updates.append(1)
+
+    yield from component.async_add_entities([entity1, entity2], True)
+
+    assert len(updates) == 1
+    assert 1 in updates