Cleanup async handling (#6388)

* Cleanups unneeded blocks

* Cleanup bootstrap

* dedicated update_ha_state

* Fix imap_email_content

* fx tests

* Fix lint & spell

homeassistant/bootstrap.py

@@ -55,6 +55,9 @@ def async_setup_component(hass: core.HomeAssistant, domain: str,
 
     This method is a coroutine.
     """
+    if domain in hass.config.components:
+        return True
+
     setup_tasks = hass.data.get(DATA_SETUP)
 
     if setup_tasks is not None and domain in setup_tasks:
@@ -211,6 +214,10 @@ def _async_setup_component(hass: core.HomeAssistant,
 
     hass.config.components.add(component.DOMAIN)
 
+    # cleanup
+    if domain in hass.data[DATA_SETUP]:
+        hass.data[DATA_SETUP].pop(domain)
+
     hass.bus.async_fire(
         EVENT_COMPONENT_LOADED, {ATTR_COMPONENT: component.DOMAIN}
     )

homeassistant/components/binary_sensor/enocean.py

@@ -67,7 +67,7 @@ class EnOceanBinarySensor(enocean.EnOceanDevice, BinarySensorDevice):
         This method is called when there is an incoming packet associated
         with this platform.
         """
-        self.update_ha_state()
+        self.schedule_update_ha_state()
         if value2 == 0x70:
             self.which = 0
             self.onoff = 0

homeassistant/components/climate/zwave.py

@@ -216,7 +216,7 @@ class ZWaveClimate(ZWaveDeviceEntity, ClimateDevice):
         self.set_value(
             class_id=zwave.const.COMMAND_CLASS_THERMOSTAT_SETPOINT,
             index=self._index, data=temperature)
-        self.update_ha_state()
+        self.schedule_update_ha_state()
 
     def set_fan_mode(self, fan):
         """Set new target fan mode."""

homeassistant/components/isy994.py

@@ -231,7 +231,7 @@ class ISYDevice(Entity):
     # pylint: disable=unused-argument
     def on_update(self, event: object) -> None:
         """Handle the update event from the ISY994 Node."""
-        self.update_ha_state()
+        self.schedule_update_ha_state()
 
     @property
     def domain(self) -> str:

homeassistant/components/light/enocean.py

@@ -106,4 +106,4 @@ class EnOceanLight(enocean.EnOceanDevice, Light):
         """Update the internal state of this device."""
         self._brightness = math.floor(val / 100.0 * 256.0)
         self._on_state = bool(val != 0)
-        self.update_ha_state()
+        self.schedule_update_ha_state()

homeassistant/components/light/lifx.py

@@ -98,7 +98,7 @@ class LIFX(object):
                           ipaddr, name, power, hue, sat, bri, kel)
             bulb.set_power(power)
             bulb.set_color(hue, sat, bri, kel)
-            bulb.update_ha_state()
+            bulb.schedule_update_ha_state()
 
     def on_color(self, ipaddr, hue, sat, bri, kel):
         """Initialize the light."""
@@ -106,7 +106,7 @@ class LIFX(object):
 
         if bulb is not None:
             bulb.set_color(hue, sat, bri, kel)
-            bulb.update_ha_state()
+            bulb.schedule_update_ha_state()
 
     def on_power(self, ipaddr, power):
         """Initialize the light."""
@@ -114,7 +114,7 @@ class LIFX(object):
 
         if bulb is not None:
             bulb.set_power(power)
-            bulb.update_ha_state()
+            bulb.schedule_update_ha_state()
 
     # pylint: disable=unused-argument
     def poll(self, now):

homeassistant/components/light/scsgate.py

@@ -106,7 +106,7 @@ class SCSGateLight(Light):
             return
 
         self._toggled = message.toggled
-        self.update_ha_state()
+        self.schedule_update_ha_state()
 
         command = "off"
         if self._toggled:

homeassistant/components/proximity.py

@@ -71,7 +71,7 @@ def setup_proximity_component(hass, name, config):
                           zone_id, unit_of_measurement)
     proximity.entity_id = '{}.{}'.format(DOMAIN, proximity_zone)
 
-    proximity.update_ha_state()
+    proximity.schedule_update_ha_state()
 
     track_state_change(
         hass, proximity_devices, proximity.check_proximity_state_change)
@@ -161,7 +161,7 @@ class Proximity(Entity):
             self.dist_to = 'not set'
             self.dir_of_travel = 'not set'
             self.nearest = 'not set'
-            self.update_ha_state()
+            self.schedule_update_ha_state()
             return
 
         # At least one device is in the monitored zone so update the entity.
@@ -169,7 +169,7 @@ class Proximity(Entity):
             self.dist_to = 0
             self.dir_of_travel = 'arrived'
             self.nearest = devices_in_zone
-            self.update_ha_state()
+            self.schedule_update_ha_state()
             return
 
         # We can't check proximity because latitude and longitude don't exist.
@@ -214,7 +214,7 @@ class Proximity(Entity):
             self.dir_of_travel = 'unknown'
             device_state = self.hass.states.get(closest_device)
             self.nearest = device_state.name
-            self.update_ha_state()
+            self.schedule_update_ha_state()
             return
 
         # Stop if we cannot calculate the direction of travel (i.e. we don't
@@ -223,7 +223,7 @@ class Proximity(Entity):
             self.dist_to = round(distances_to_zone[entity])
             self.dir_of_travel = 'unknown'
             self.nearest = entity_name
-            self.update_ha_state()
+            self.schedule_update_ha_state()
             return
 
         # Reset the variables
@@ -250,7 +250,7 @@ class Proximity(Entity):
         self.dist_to = round(dist_to_zone)
         self.dir_of_travel = direction_of_travel
         self.nearest = entity_name
-        self.update_ha_state()
+        self.schedule_update_ha_state()
         _LOGGER.debug('proximity.%s update entity: distance=%s: direction=%s: '
                       'device=%s', self.friendly_name, round(dist_to_zone),
                       direction_of_travel, entity_name)

homeassistant/components/qwikswitch.py

@@ -89,7 +89,7 @@ class QSToggleEntity(object):
         if value != self._value:
             self._value = value
             # pylint: disable=no-member
-            super().update_ha_state()  # Part of Entity/ToggleEntity
+            super().schedule_update_ha_state()  # Part of Entity/ToggleEntity
         return self._value
 
     def turn_on(self, **kwargs):

homeassistant/components/remote/harmony.py

@@ -92,7 +92,7 @@ def _apply_service(service, service_func, *service_func_args):
 
     for device in _devices:
         service_func(device, *service_func_args)
-        device.update_ha_state(True)
+        device.schedule_update_ha_state(True)
 
 
 def _sync_service(service):

homeassistant/components/rfxtrx.py

@@ -334,7 +334,7 @@ class RfxtrxDevice(Entity):
         """Update det state of the device."""
         self._state = state
         self._brightness = brightness
-        self.update_ha_state()
+        self.schedule_update_ha_state()
 
     def _send_command(self, command, brightness=0):
         if not self._event:
@@ -369,4 +369,4 @@ class RfxtrxDevice(Entity):
             for _ in range(self.signal_repetitions):
                 self._event.device.send_stop(RFXOBJECT.transport)
 
-        self.update_ha_state()
+        self.schedule_update_ha_state()

homeassistant/components/sensor/enocean.py

@@ -54,7 +54,7 @@ class EnOceanSensor(enocean.EnOceanDevice, Entity):
     def value_changed(self, value):
         """Update the internal state of the device."""
         self.power = value
-        self.update_ha_state()
+        self.schedule_update_ha_state()
 
     @property
     def state(self):

homeassistant/components/sensor/fritzbox_callmonitor.py

@@ -169,4 +169,4 @@ class FritzBoxCallMonitor(object):
             self._sensor.set_state(VALUE_DISCONNECT)
             att = {"duration": line[3], "closed": isotime}
             self._sensor.set_attributes(att)
-        self._sensor.update_ha_state()
+        self._sensor.schedule_update_ha_state()

homeassistant/components/sensor/haveibeenpwned.py

@@ -110,7 +110,7 @@ class HaveIBeenPwnedSensor(Entity):
 
         if self._email in self._data.data:
             self._state = len(self._data.data[self._email])
-            self.update_ha_state()
+            self.schedule_update_ha_state()
 
     def update(self):
         """Update data and see if it contains data for our email."""

homeassistant/components/sensor/imap_email_content.py

@@ -225,25 +225,23 @@ class EmailContentSensor(Entity):
 
     def update(self):
         """Read emails and publish state change."""
-        while True:
+        email_message = self._email_reader.read_next()
-            email_message = self._email_reader.read_next()
 
-            if email_message is None:
+        if email_message is None:
-                break
+            return
 
-            if self.sender_allowed(email_message):
+        if self.sender_allowed(email_message):
-                message_body = EmailContentSensor.get_msg_text(email_message)
+            message_body = EmailContentSensor.get_msg_text(email_message)
 
-                if self._value_template is not None:
+            if self._value_template is not None:
-                    message_body = self.render_template(email_message)
+                message_body = self.render_template(email_message)
 
-                self._message = message_body
+            self._message = message_body
-                self._state_attributes = {
+            self._state_attributes = {
-                    ATTR_FROM:
+                ATTR_FROM:
-                        EmailContentSensor.get_msg_sender(email_message),
+                    EmailContentSensor.get_msg_sender(email_message),
-                    ATTR_SUBJECT:
+                ATTR_SUBJECT:
-                        EmailContentSensor.get_msg_subject(email_message),
+                    EmailContentSensor.get_msg_subject(email_message),
-                    ATTR_DATE:
+                ATTR_DATE:
-                        email_message['Date']
+                    email_message['Date']
-                }
+            }
-                self.update_ha_state()

homeassistant/components/sensor/loopenergy.py

@@ -121,7 +121,7 @@ class LoopEnergyDevice(Entity):
         return self._unit_of_measurement
 
     def _callback(self):
-        self.update_ha_state(True)
+        self.schedule_update_ha_state(True)
 
 
 class LoopEnergyElec(LoopEnergyDevice):

homeassistant/components/sensor/pilight.py

@@ -89,7 +89,7 @@ class PilightSensor(Entity):
             try:
                 value = call.data[self._variable]
                 self._state = value
-                self.update_ha_state()
+                self.schedule_update_ha_state()
             except KeyError:
                 _LOGGER.error(
                     'No variable %s in received code data %s',

homeassistant/components/sun.py

@@ -231,7 +231,7 @@ class Sun(Entity):
     def point_in_time_listener(self, now):
         """Called when the state of the sun has changed."""
         self.update_as_of(now)
-        self.update_ha_state()
+        self.schedule_update_ha_state()
 
         # Schedule next update at next_change+1 second so sun state has changed
         track_point_in_utc_time(
@@ -241,4 +241,4 @@ class Sun(Entity):
     def timer_update(self, time):
         """Needed to update solar elevation and azimuth."""
         self.update_sun_position(time)
-        self.update_ha_state()
+        self.schedule_update_ha_state()

homeassistant/components/switch/broadlink.py

@@ -188,13 +188,13 @@ class BroadlinkRMSwitch(SwitchDevice):
         """Turn the device on."""
         if self._sendpacket(self._command_on):
             self._state = True
-            self.update_ha_state()
+            self.schedule_update_ha_state()
 
     def turn_off(self, **kwargs):
         """Turn the device off."""
         if self._sendpacket(self._command_off):
             self._state = False
-            self.update_ha_state()
+            self.schedule_update_ha_state()
 
     def _sendpacket(self, packet, retry=2):
         """Send packet to device."""

homeassistant/components/switch/rpi_rf.py

@@ -114,10 +114,10 @@ class RPiRFSwitch(SwitchDevice):
         """Turn the switch on."""
         if self._send_code(self._code_on, self._protocol, self._pulselength):
             self._state = True
-            self.update_ha_state()
+            self.schedule_update_ha_state()
 
     def turn_off(self):
         """Turn the switch off."""
         if self._send_code(self._code_off, self._protocol, self._pulselength):
             self._state = False
-            self.update_ha_state()
+            self.schedule_update_ha_state()

homeassistant/components/switch/scsgate.py

@@ -141,7 +141,7 @@ class SCSGateSwitch(SwitchDevice):
             return
 
         self._toggled = message.toggled
-        self.update_ha_state()
+        self.schedule_update_ha_state()
 
         command = "off"
         if self._toggled:

homeassistant/components/weblink.py

@@ -54,7 +54,7 @@ class Link(Entity):
         self._url = url
         self._icon = icon
         self.entity_id = DOMAIN + '.%s' % slugify(name)
-        self.update_ha_state()
+        self.schedule_update_ha_state()
 
     @property
     def icon(self):

homeassistant/components/wink.py

@@ -115,7 +115,7 @@ def setup(hass, config):
         """Force all devices to poll the Wink API."""
         _LOGGER.info("Refreshing Wink states from API")
         for entity in hass.data[DOMAIN]['entities']:
-            entity.update_ha_state(True)
+            entity.schedule_update_ha_state(True)
     hass.services.register(DOMAIN, 'Refresh state from Wink', force_update)
 
     def pull_new_devices(call):
@@ -150,14 +150,14 @@ class WinkDevice(Entity):
             if message is None:
                 _LOGGER.error("Error on pubnub update for %s "
                               "polling API for current state", self.name)
-                self.update_ha_state(True)
+                self.schedule_update_ha_state(True)
             else:
                 self.wink.pubnub_update(message)
-                self.update_ha_state()
+                self.schedule_update_ha_state()
         except (ValueError, KeyError, AttributeError):
             _LOGGER.error("Error in pubnub JSON for %s "
                           "polling API for current state", self.name)
-            self.update_ha_state(True)
+            self.schedule_update_ha_state(True)
 
     @property
     def name(self):

homeassistant/helpers/entity.py

@@ -180,14 +180,9 @@ class Entity(object):
 
         If force_refresh == True will update entity before setting state.
         """
-        # We're already in a thread, do the force refresh here.
+        _LOGGER.warning("'update_ha_state' is deprecated. "
-        if force_refresh and not hasattr(self, 'async_update'):
+                        "Use 'schedule_update_ha_state' instead.")
-            self.update()
+        self.schedule_update_ha_state(force_refresh)
-            force_refresh = False
-
-        run_coroutine_threadsafe(
-            self.async_update_ha_state(force_refresh), self.hass.loop
-        ).result()
 
     @asyncio.coroutine
     def async_update_ha_state(self, force_refresh=False):
@@ -280,11 +275,6 @@ class Entity(object):
 
         That is only needed on executor to not block.
         """
-        # We're already in a thread, do the force refresh here.
-        if force_refresh and not hasattr(self, 'async_update'):
-            self.update()
-            force_refresh = False
-
         self.hass.add_job(self.async_update_ha_state(force_refresh))
 
     def remove(self) -> None:

homeassistant/helpers/entity_component.py

@@ -309,13 +309,10 @@ class EntityPlatform(object):
 
     def schedule_add_entities(self, new_entities, update_before_add=False):
         """Add entities for a single platform."""
-        if update_before_add:
-            for entity in new_entities:
-                entity.update()
-
         run_callback_threadsafe(
             self.component.hass.loop,
-            self.async_schedule_add_entities, list(new_entities), False
+            self.async_schedule_add_entities, list(new_entities),
+            update_before_add
         ).result()
 
     @callback

tests/components/media_player/test_universal.py

@@ -140,10 +140,12 @@ class TestMediaPlayer(unittest.TestCase):
         self.hass = get_test_home_assistant()
 
         self.mock_mp_1 = MockMediaPlayer(self.hass, 'mock1')
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
 
         self.mock_mp_2 = MockMediaPlayer(self.hass, 'mock2')
-        self.mock_mp_2.update_ha_state()
+        self.mock_mp_2.schedule_update_ha_state()
+
+        self.hass.block_till_done()
 
         self.mock_mute_switch_id = switch.ENTITY_ID_FORMAT.format('mute')
         self.hass.states.set(self.mock_mute_switch_id, STATE_OFF)
@@ -315,19 +317,22 @@ class TestMediaPlayer(unittest.TestCase):
         self.assertEqual(None, ump._child_state)
 
         self.mock_mp_1._state = STATE_PLAYING
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertEqual(self.mock_mp_1.entity_id,
                          ump._child_state.entity_id)
 
         self.mock_mp_2._state = STATE_PLAYING
-        self.mock_mp_2.update_ha_state()
+        self.mock_mp_2.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertEqual(self.mock_mp_1.entity_id,
                          ump._child_state.entity_id)
 
         self.mock_mp_1._state = STATE_OFF
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertEqual(self.mock_mp_2.entity_id,
                          ump._child_state.entity_id)
@@ -362,7 +367,8 @@ class TestMediaPlayer(unittest.TestCase):
         self.assertTrue(ump.state, STATE_OFF)
 
         self.mock_mp_1._state = STATE_PLAYING
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertEqual(STATE_PLAYING, ump.state)
 
@@ -382,7 +388,8 @@ class TestMediaPlayer(unittest.TestCase):
         self.assertEqual(STATE_ON, ump.state)
 
         self.mock_mp_1._state = STATE_PLAYING
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertEqual(STATE_PLAYING, ump.state)
 
@@ -402,12 +409,14 @@ class TestMediaPlayer(unittest.TestCase):
         self.assertEqual(None, ump.volume_level)
 
         self.mock_mp_1._state = STATE_PLAYING
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertEqual(0, ump.volume_level)
 
         self.mock_mp_1._volume_level = 1
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertEqual(1, ump.volume_level)
 
@@ -425,7 +434,8 @@ class TestMediaPlayer(unittest.TestCase):
 
         self.mock_mp_1._state = STATE_PLAYING
         self.mock_mp_1._media_image_url = TEST_URL
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         # mock_mp_1 will convert the url to the api proxy url. This test
         # ensures ump passes through the same url without an additional proxy.
@@ -443,12 +453,14 @@ class TestMediaPlayer(unittest.TestCase):
         self.assertFalse(ump.is_volume_muted)
 
         self.mock_mp_1._state = STATE_PLAYING
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertFalse(ump.is_volume_muted)
 
         self.mock_mp_1._is_volume_muted = True
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertTrue(ump.is_volume_muted)
 
@@ -513,7 +525,8 @@ class TestMediaPlayer(unittest.TestCase):
 
         self.mock_mp_1._supported_features = 512
         self.mock_mp_1._state = STATE_PLAYING
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
         self.assertEqual(512, ump.supported_features)
 
@@ -534,7 +547,8 @@ class TestMediaPlayer(unittest.TestCase):
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
 
         self.mock_mp_1._state = STATE_PLAYING
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
 
         check_flags = universal.SUPPORT_TURN_ON | universal.SUPPORT_TURN_OFF \
@@ -553,9 +567,10 @@ class TestMediaPlayer(unittest.TestCase):
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
 
         self.mock_mp_1._state = STATE_OFF
-        self.mock_mp_1.update_ha_state()
+        self.mock_mp_1.schedule_update_ha_state()
         self.mock_mp_2._state = STATE_OFF
-        self.mock_mp_2.update_ha_state()
+        self.mock_mp_2.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
 
         run_coroutine_threadsafe(
@@ -574,7 +589,8 @@ class TestMediaPlayer(unittest.TestCase):
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
 
         self.mock_mp_2._state = STATE_PLAYING
-        self.mock_mp_2.update_ha_state()
+        self.mock_mp_2.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
 
         run_coroutine_threadsafe(
@@ -672,7 +688,8 @@ class TestMediaPlayer(unittest.TestCase):
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
 
         self.mock_mp_2._state = STATE_PLAYING
-        self.mock_mp_2.update_ha_state()
+        self.mock_mp_2.schedule_update_ha_state()
+        self.hass.block_till_done()
         run_coroutine_threadsafe(ump.async_update(), self.hass.loop).result()
 
         run_coroutine_threadsafe(ump.async_turn_off(), self.hass.loop).result()

tests/components/sensor/test_imap_email_content.py

@@ -4,7 +4,6 @@ import email
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
 import datetime
-from threading import Event
 import unittest
 
 from homeassistant.helpers.template import Template
@@ -59,7 +58,8 @@ class EmailContentSensor(unittest.TestCase):
             None)
 
         sensor.entity_id = 'sensor.emailtest'
-        sensor.update()
+        sensor.schedule_update_ha_state(True)
+        self.hass.block_till_done()
         self.assertEqual("Test Message", sensor.state)
         self.assertEqual('sender@test.com',
                          sensor.device_state_attributes['from'])
@@ -87,7 +87,8 @@ class EmailContentSensor(unittest.TestCase):
             ['sender@test.com'], None)
 
         sensor.entity_id = "sensor.emailtest"
-        sensor.update()
+        sensor.schedule_update_ha_state(True)
+        self.hass.block_till_done()
         self.assertEqual("Test Message", sensor.state)
 
     def test_multi_part_only_html(self):
@@ -110,7 +111,8 @@ class EmailContentSensor(unittest.TestCase):
             None)
 
         sensor.entity_id = 'sensor.emailtest'
-        sensor.update()
+        sensor.schedule_update_ha_state(True)
+        self.hass.block_till_done()
         self.assertEqual(
             "<html><head></head><body>Test Message</body></html>",
             sensor.state)
@@ -132,7 +134,8 @@ class EmailContentSensor(unittest.TestCase):
             ['sender@test.com'], None)
 
         sensor.entity_id = 'sensor.emailtest'
-        sensor.update()
+        sensor.schedule_update_ha_state(True)
+        self.hass.block_till_done()
         self.assertEqual("Test Message", sensor.state)
 
     def test_multiple_emails(self):
@@ -151,12 +154,8 @@ class EmailContentSensor(unittest.TestCase):
         test_message2['Date'] = datetime.datetime(2016, 1, 1, 12, 44, 57)
         test_message2.set_payload("Test Message 2")
 
-        states_received = Event()
-
         def state_changed_listener(entity_id, from_s, to_s):
             states.append(to_s)
-            if len(states) == 2:
-                states_received.set()
 
         track_state_change(
             self.hass, ['sensor.emailtest'], state_changed_listener)
@@ -167,10 +166,11 @@ class EmailContentSensor(unittest.TestCase):
             'test_emails_sensor', ['sender@test.com'], None)
 
         sensor.entity_id = 'sensor.emailtest'
-        sensor.update()
 
+        sensor.schedule_update_ha_state(True)
+        self.hass.block_till_done()
+        sensor.schedule_update_ha_state(True)
         self.hass.block_till_done()
-        states_received.wait(5)
 
         self.assertEqual("Test Message", states[0].state)
         self.assertEqual("Test Message 2", states[1].state)
@@ -190,7 +190,8 @@ class EmailContentSensor(unittest.TestCase):
             'test_emails_sensor', ['other@test.com'], None)
 
         sensor.entity_id = 'sensor.emailtest'
-        sensor.update()
+        sensor.schedule_update_ha_state(True)
+        self.hass.block_till_done()
         self.assertEqual(None, sensor.state)
 
     def test_template(self):
@@ -208,7 +209,8 @@ class EmailContentSensor(unittest.TestCase):
                      self.hass))
 
         sensor.entity_id = 'sensor.emailtest'
-        sensor.update()
+        sensor.schedule_update_ha_state(True)
+        self.hass.block_till_done()
         self.assertEqual(
             "Test from sender@test.com with message Test Message",
             sensor.state)

tests/components/test_init.py

@@ -103,7 +103,8 @@ class TestComponentsCore(unittest.TestCase):
         ent = entity.Entity()
         ent.entity_id = 'test.entity'
         ent.hass = self.hass
-        ent.update_ha_state()
+        ent.schedule_update_ha_state()
+        self.hass.block_till_done()
 
         state = self.hass.states.get('test.entity')
         assert state is not None
@@ -130,7 +131,8 @@ class TestComponentsCore(unittest.TestCase):
         assert 10 == self.hass.config.latitude
         assert 20 == self.hass.config.longitude
 
-        ent.update_ha_state()
+        ent.schedule_update_ha_state()
+        self.hass.block_till_done()
 
         state = self.hass.states.get('test.entity')
         assert state is not None

tests/helpers/test_entity.py

@@ -77,7 +77,8 @@ class TestHelpersEntity(object):
         self.entity = entity.Entity()
         self.entity.entity_id = 'test.overwrite_hidden_true'
         self.hass = self.entity.hass = get_test_home_assistant()
-        self.entity.update_ha_state()
+        self.entity.schedule_update_ha_state()
+        self.hass.block_till_done()
 
     def teardown_method(self, method):
         """Stop everything that was started."""
@@ -92,7 +93,8 @@ class TestHelpersEntity(object):
         """Test we can overwrite hidden property to True."""
         self.hass.data[DATA_CUSTOMIZE] = EntityValues({
             self.entity.entity_id: {ATTR_HIDDEN: True}})
-        self.entity.update_ha_state()
+        self.entity.schedule_update_ha_state()
+        self.hass.block_till_done()
 
         state = self.hass.states.get(self.entity.entity_id)
         assert state.attributes.get(ATTR_HIDDEN)
@@ -126,6 +128,7 @@ class TestHelpersEntity(object):
         assert state.attributes.get(ATTR_DEVICE_CLASS) is None
         with patch('homeassistant.helpers.entity.Entity.device_class',
                    new='test_class'):
-            self.entity.update_ha_state()
+            self.entity.schedule_update_ha_state()
+            self.hass.block_till_done()
         state = self.hass.states.get(self.entity.entity_id)
         assert state.attributes.get(ATTR_DEVICE_CLASS) == 'test_class'

tests/test_config.py

@@ -213,7 +213,7 @@ class TestConfig(unittest.TestCase):
         entity = Entity()
         entity.entity_id = 'test.test'
         entity.hass = self.hass
-        entity.update_ha_state()
+        entity.schedule_update_ha_state()
 
         self.hass.block_till_done()