Fix automations listening to HOMEASSISTANT_START (#6936)

* Fire EVENT_HOMEASSISTANT_START automations off right away while starting

* Actually have core state be set to 'starting' during boot

* Fix correct start implementation

* Test and deprecate event automation platform on start

* Fix doc strings

* Remove shutting down exception

* More strict when to mark an instance as finished

* Add automation platform to listen for start/shutdown

* When we stop we should wait till it's all done

* Fix testing

* Fix async bugs in tests

* Only set UVLOOP when hass starts from CLI

* This hangs normal asyncio event loop

* Clean up Z-Wave node entity test

homeassistant/__main__.py

@@ -20,6 +20,17 @@ from homeassistant.const import (
 from homeassistant.util.async import run_callback_threadsafe
 
 
+def attempt_use_uvloop():
+    """Attempt to use uvloop."""
+    import asyncio
+
+    try:
+        import uvloop
+        asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+    except ImportError:
+        pass
+
+
 def monkey_patch_asyncio():
     """Replace weakref.WeakSet to address Python 3 bug.
 
@@ -311,8 +322,7 @@ def setup_and_run_hass(config_dir: str,
             EVENT_HOMEASSISTANT_START, open_browser
         )
 
-    hass.start()
-    return hass.exit_code
+    return hass.start()
 
 
 def try_to_restart() -> None:
@@ -359,11 +369,13 @@ def try_to_restart() -> None:
 
 def main() -> int:
     """Start Home Assistant."""
+    validate_python()
+
+    attempt_use_uvloop()
+
     if sys.version_info[:3] < (3, 5, 3):
         monkey_patch_asyncio()
 
-    validate_python()
-
     args = get_arguments()
 
     if args.script is not None:

homeassistant/bootstrap.py

@@ -74,8 +74,6 @@ def async_from_config_dict(config: Dict[str, Any],
     This method is a coroutine.
     """
     start = time()
-    hass.async_track_tasks()
-
     core_config = config.get(core.DOMAIN, {})
 
     try:
@@ -140,10 +138,10 @@ def async_from_config_dict(config: Dict[str, Any],
             continue
         hass.async_add_job(async_setup_component(hass, component, config))
 
-    yield from hass.async_stop_track_tasks()
+    yield from hass.async_block_till_done()
 
     stop = time()
-    _LOGGER.info('Home Assistant initialized in %ss', round(stop-start, 2))
+    _LOGGER.info('Home Assistant initialized in %.2fs', stop-start)
 
     async_register_signal_handling(hass)
     return hass

homeassistant/components/automation/event.py

@@ -2,15 +2,15 @@
 Offer event listening automation rules.
 
 For more details about this automation rule, please refer to the documentation
-at https://home-assistant.io/components/automation/#event-trigger
+at https://home-assistant.io/docs/automation/trigger/#event-trigger
 """
 import asyncio
 import logging
 
 import voluptuous as vol
 
-from homeassistant.core import callback
-from homeassistant.const import CONF_PLATFORM
+from homeassistant.core import callback, CoreState
+from homeassistant.const import CONF_PLATFORM, EVENT_HOMEASSISTANT_START
 from homeassistant.helpers import config_validation as cv
 
 CONF_EVENT_TYPE = "event_type"
@@ -31,6 +31,19 @@ def async_trigger(hass, config, action):
     event_type = config.get(CONF_EVENT_TYPE)
     event_data = config.get(CONF_EVENT_DATA)
 
+    if (event_type == EVENT_HOMEASSISTANT_START and
+            hass.state == CoreState.starting):
+        _LOGGER.warning('Deprecation: Automations should not listen to event '
+                        "'homeassistant_start'. Use platform 'homeassistant' "
+                        'instead. Feature will be removed in 0.45')
+        hass.async_run_job(action, {
+            'trigger': {
+                'platform': 'event',
+                'event': None,
+            },
+        })
+        return lambda: None
+
     @callback
     def handle_event(event):
         """Listen for events and calls the action when data matches."""

homeassistant/components/automation/homeassistant.py

@@ -0,0 +1,55 @@
+"""
+Offer Home Assistant core automation rules.
+
+For more details about this automation rule, please refer to the documentation
+at https://home-assistant.io/components/automation/#homeassistant-trigger
+"""
+import asyncio
+import logging
+
+import voluptuous as vol
+
+from homeassistant.core import callback, CoreState
+from homeassistant.const import (
+    CONF_PLATFORM, CONF_EVENT, EVENT_HOMEASSISTANT_STOP)
+
+EVENT_START = 'start'
+EVENT_SHUTDOWN = 'shutdown'
+_LOGGER = logging.getLogger(__name__)
+
+TRIGGER_SCHEMA = vol.Schema({
+    vol.Required(CONF_PLATFORM): 'homeassistant',
+    vol.Required(CONF_EVENT): vol.Any(EVENT_START, EVENT_SHUTDOWN),
+})
+
+
+@asyncio.coroutine
+def async_trigger(hass, config, action):
+    """Listen for events based on configuration."""
+    event = config.get(CONF_EVENT)
+
+    if event == EVENT_SHUTDOWN:
+        @callback
+        def hass_shutdown(event):
+            """Called when Home Assistant is shutting down."""
+            hass.async_run_job(action, {
+                'trigger': {
+                    'platform': 'homeassistant',
+                    'event': event,
+                },
+            })
+
+        return hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP,
+                                          hass_shutdown)
+
+    # Automation are enabled while hass is starting up, fire right away
+    # Check state because a config reload shouldn't trigger it.
+    elif hass.state == CoreState.starting:
+        hass.async_run_job(action, {
+            'trigger': {
+                'platform': 'homeassistant',
+                'event': event,
+            },
+        })
+
+    return lambda: None

homeassistant/components/automation/litejet.py

@@ -70,7 +70,7 @@ def async_trigger(hass, config, action):
         nonlocal held_less_than, held_more_than
         pressed_time = dt_util.utcnow()
         if held_more_than is None and held_less_than is None:
-            call_action()
+            hass.add_job(call_action)
         if held_more_than is not None and held_less_than is None:
             cancel_pressed_more_than = track_point_in_utc_time(
                 hass,
@@ -88,7 +88,7 @@ def async_trigger(hass, config, action):
         held_time = dt_util.utcnow() - pressed_time
         if held_less_than is not None and held_time < held_less_than:
             if held_more_than is None or held_time > held_more_than:
-                call_action()
+                hass.add_job(call_action)
 
     hass.data['litejet_system'].on_switch_pressed(number, pressed)
     hass.data['litejet_system'].on_switch_released(number, released)

homeassistant/components/automation/mqtt.py

@@ -2,7 +2,7 @@
 Offer MQTT listening automation rules.
 
 For more details about this automation rule, please refer to the documentation
-at https://home-assistant.io/components/automation/#mqtt-trigger
+at https://home-assistant.io/docs/automation/trigger/#mqtt-trigger
 """
 import asyncio
 import json

homeassistant/components/automation/numeric_state.py

@@ -2,7 +2,7 @@
 Offer numeric state listening automation rules.
 
 For more details about this automation rule, please refer to the documentation
-at https://home-assistant.io/components/automation/#numeric-state-trigger
+at https://home-assistant.io/docs/automation/trigger/#numeric-state-trigger
 """
 import asyncio
 import logging

homeassistant/components/automation/state.py

@@ -2,7 +2,7 @@
 Offer state listening automation rules.
 
 For more details about this automation rule, please refer to the documentation
-at https://home-assistant.io/components/automation/#state-trigger
+at https://home-assistant.io/docs/automation/trigger/#state-trigger
 """
 import asyncio
 import voluptuous as vol

homeassistant/components/automation/sun.py

@@ -2,7 +2,7 @@
 Offer sun based automation rules.
 
 For more details about this automation rule, please refer to the documentation
-at https://home-assistant.io/components/automation/#sun-trigger
+at https://home-assistant.io/docs/automation/trigger/#sun-trigger
 """
 import asyncio
 from datetime import timedelta

homeassistant/components/automation/template.py

@@ -2,7 +2,7 @@
 Offer template automation rules.
 
 For more details about this automation rule, please refer to the documentation
-at https://home-assistant.io/components/automation/#template-trigger
+at https://home-assistant.io/docs/automation/trigger/#template-trigger
 """
 import asyncio
 import logging

homeassistant/components/automation/time.py

@@ -2,7 +2,7 @@
 Offer time listening automation rules.
 
 For more details about this automation rule, please refer to the documentation
-at https://home-assistant.io/components/automation/#time-trigger
+at https://home-assistant.io/docs/automation/trigger/#time-trigger
 """
 import asyncio
 import logging

homeassistant/components/automation/zone.py

@@ -2,7 +2,7 @@
 Offer zone automation rules.
 
 For more details about this automation rule, please refer to the documentation
-at https://home-assistant.io/components/automation/#zone-trigger
+at https://home-assistant.io/docs/automation/trigger/#zone-trigger
 """
 import asyncio
 import voluptuous as vol

homeassistant/components/light/litejet.py

@@ -48,7 +48,7 @@ class LiteJetLight(Light):
     def _on_load_changed(self):
         """Called on a LiteJet thread when a load's state changes."""
         _LOGGER.debug("Updating due to notification for %s", self._name)
-        self._hass.async_add_job(self.async_update_ha_state(True))
+        self.schedule_update_ha_state(True)
 
     @property
     def supported_features(self):

homeassistant/components/sensor/mqtt_room.py

@@ -98,6 +98,7 @@ class MQTTRoomSensor(Entity):
 
             self.hass.async_add_job(self.async_update_ha_state())
 
+        @callback
         def message_received(topic, payload, qos):
             """A new MQTT message has been received."""
             try:

homeassistant/components/switch/litejet.py

@@ -47,12 +47,12 @@ class LiteJetSwitch(SwitchDevice):
     def _on_switch_pressed(self):
         _LOGGER.debug("Updating pressed for %s", self._name)
         self._state = True
-        self._hass.async_add_job(self.async_update_ha_state())
+        self.schedule_update_ha_state()
 
     def _on_switch_released(self):
         _LOGGER.debug("Updating released for %s", self._name)
         self._state = False
-        self._hass.async_add_job(self.async_update_ha_state())
+        self.schedule_update_ha_state()
 
     @property
     def name(self):

homeassistant/core.py

@@ -29,7 +29,7 @@ from homeassistant.const import (
     EVENT_TIME_CHANGED, MATCH_ALL, EVENT_HOMEASSISTANT_CLOSE,
     EVENT_SERVICE_REMOVED, __version__)
 from homeassistant.exceptions import (
-    HomeAssistantError, InvalidEntityFormatError, ShuttingDown)
+    HomeAssistantError, InvalidEntityFormatError)
 from homeassistant.util.async import (
     run_coroutine_threadsafe, run_callback_threadsafe)
 import homeassistant.util as util
@@ -37,12 +37,6 @@ import homeassistant.util.dt as dt_util
 import homeassistant.util.location as location
 from homeassistant.util.unit_system import UnitSystem, METRIC_SYSTEM  # NOQA
 
-try:
-    import uvloop
-    asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
-except ImportError:
-    pass
-
 DOMAIN = 'homeassistant'
 
 # How long we wait for the result of a service call
@@ -86,10 +80,6 @@ def async_loop_exception_handler(loop, context):
     kwargs = {}
     exception = context.get('exception')
     if exception:
-        # Do not report on shutting down exceptions.
-        if isinstance(exception, ShuttingDown):
-            return
-
         kwargs['exc_info'] = (type(exception), exception,
                               exception.__traceback__)
 
@@ -123,7 +113,7 @@ class HomeAssistant(object):
         self.loop.set_default_executor(self.executor)
         self.loop.set_exception_handler(async_loop_exception_handler)
         self._pending_tasks = []
-        self._track_task = False
+        self._track_task = True
         self.bus = EventBus(self)
         self.services = ServiceRegistry(self)
         self.states = StateMachine(self.bus, self.loop)
@@ -148,6 +138,7 @@ class HomeAssistant(object):
             # Block until stopped
             _LOGGER.info("Starting Home Assistant core loop")
             self.loop.run_forever()
+            return self.exit_code
         except KeyboardInterrupt:
             self.loop.create_task(self.async_stop())
             self.loop.run_forever()
@@ -165,9 +156,10 @@ class HomeAssistant(object):
 
         # pylint: disable=protected-access
         self.loop._thread_ident = threading.get_ident()
-        _async_create_timer(self)
         self.bus.async_fire(EVENT_HOMEASSISTANT_START)
+        yield from self.async_stop_track_tasks()
         self.state = CoreState.running
+        _async_create_timer(self)
 
     def add_job(self, target: Callable[..., None], *args: Any) -> None:
         """Add job to the executor pool.
@@ -238,6 +230,8 @@ class HomeAssistant(object):
     @asyncio.coroutine
     def async_block_till_done(self):
         """Block till all pending work is done."""
+        assert self._track_task, 'Not tracking tasks'
+
         # To flush out any call_soon_threadsafe
         yield from asyncio.sleep(0, loop=self.loop)
 
@@ -252,7 +246,8 @@ class HomeAssistant(object):
 
     def stop(self) -> None:
         """Stop Home Assistant and shuts down all threads."""
-        run_coroutine_threadsafe(self.async_stop(), self.loop)
+        self.loop.call_soon_threadsafe(
+            self.loop.create_task, self.async_stop())
 
     @asyncio.coroutine
     def async_stop(self, exit_code=0) -> None:
@@ -368,10 +363,6 @@ class EventBus(object):
 
         This method must be run in the event loop.
         """
-        if event_type != EVENT_HOMEASSISTANT_STOP and \
-                self._hass.state == CoreState.stopping:
-            raise ShuttingDown("Home Assistant is shutting down")
-
         listeners = self._listeners.get(event_type, [])
 
         # EVENT_HOMEASSISTANT_CLOSE should go only to his listeners

homeassistant/exceptions.py

@@ -7,12 +7,6 @@ class HomeAssistantError(Exception):
     pass
 
 
-class ShuttingDown(HomeAssistantError):
-    """When trying to change something during shutdown."""
-
-    pass
-
-
 class InvalidEntityFormatError(HomeAssistantError):
     """When an invalid formatted entity is encountered."""
 

tests/common.py

@@ -23,12 +23,13 @@ import homeassistant.util.yaml as yaml
 from homeassistant.const import (
     STATE_ON, STATE_OFF, DEVICE_DEFAULT_NAME, EVENT_TIME_CHANGED,
     EVENT_STATE_CHANGED, EVENT_PLATFORM_DISCOVERED, ATTR_SERVICE,
-    ATTR_DISCOVERED, SERVER_PORT, EVENT_HOMEASSISTANT_STOP)
+    ATTR_DISCOVERED, SERVER_PORT, EVENT_HOMEASSISTANT_CLOSE)
 from homeassistant.components import sun, mqtt, recorder
 from homeassistant.components.http.auth import auth_middleware
 from homeassistant.components.http.const import (
     KEY_USE_X_FORWARDED_FOR, KEY_BANS_ENABLED, KEY_TRUSTED_NETWORKS)
-from homeassistant.util.async import run_callback_threadsafe
+from homeassistant.util.async import (
+    run_callback_threadsafe, run_coroutine_threadsafe)
 
 _TEST_INSTANCE_PORT = SERVER_PORT
 _LOGGER = logging.getLogger(__name__)
@@ -58,15 +59,11 @@ def get_test_home_assistant():
         loop.run_forever()
         stop_event.set()
 
-    orig_start = hass.start
     orig_stop = hass.stop
 
-    @patch.object(hass.loop, 'run_forever')
-    @patch.object(hass.loop, 'close')
     def start_hass(*mocks):
         """Helper to start hass."""
-        orig_start()
-        hass.block_till_done()
+        run_coroutine_threadsafe(hass.async_start(), loop=hass.loop).result()
 
     def stop_hass():
         """Stop hass."""
@@ -101,7 +98,6 @@ def async_test_home_assistant(loop):
         return orig_async_add_job(target, *args)
 
     hass.async_add_job = async_add_job
-    hass.async_track_tasks()
 
     hass.config.location_name = 'test home'
     hass.config.config_dir = get_test_config_dir()
@@ -123,7 +119,11 @@ def async_test_home_assistant(loop):
     @asyncio.coroutine
     def mock_async_start():
         """Start the mocking."""
-        with patch('homeassistant.core._async_create_timer'):
+        # 1. We only mock time during tests
+        # 2. We want block_till_done that is called inside stop_track_tasks
+        with patch('homeassistant.core._async_create_timer'), \
+                patch.object(hass, 'async_stop_track_tasks',
+                             hass.async_block_till_done):
             yield from orig_start()
 
     hass.async_start = mock_async_start
@@ -134,7 +134,7 @@ def async_test_home_assistant(loop):
         global INST_COUNT
         INST_COUNT -= 1
 
-    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, clear_instance)
+    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_CLOSE, clear_instance)
 
     return hass
 

tests/components/automation/test_event.py

@@ -1,11 +1,13 @@
 """The tests for the Event automation."""
+import asyncio
 import unittest
 
-from homeassistant.core import callback
-from homeassistant.setup import setup_component
+from homeassistant.const import EVENT_HOMEASSISTANT_START
+from homeassistant.core import callback, CoreState
+from homeassistant.setup import setup_component, async_setup_component
 import homeassistant.components.automation as automation
 
-from tests.common import get_test_home_assistant, mock_component
+from tests.common import get_test_home_assistant, mock_component, mock_service
 
 
 # pylint: disable=invalid-name
@@ -92,3 +94,30 @@ class TestAutomationEvent(unittest.TestCase):
         self.hass.bus.fire('test_event', {'some_attr': 'some_other_value'})
         self.hass.block_till_done()
         self.assertEqual(0, len(self.calls))
+
+
+@asyncio.coroutine
+def test_if_fires_on_event_with_data(hass):
+    """Test the firing of events with data."""
+    calls = mock_service(hass, 'test', 'automation')
+    hass.state = CoreState.not_running
+
+    res = yield from async_setup_component(hass, automation.DOMAIN, {
+        automation.DOMAIN: {
+            'alias': 'hello',
+            'trigger': {
+                'platform': 'event',
+                'event_type': EVENT_HOMEASSISTANT_START,
+            },
+            'action': {
+                'service': 'test.automation',
+            }
+        }
+    })
+    assert res
+    assert not automation.is_on(hass, 'automation.hello')
+    assert len(calls) == 0
+
+    yield from hass.async_start()
+    assert automation.is_on(hass, 'automation.hello')
+    assert len(calls) == 1

tests/components/automation/test_homeassistant.py

@@ -0,0 +1,84 @@
+"""The tests for the Event automation."""
+import asyncio
+from unittest.mock import patch, Mock
+
+from homeassistant.core import CoreState
+from homeassistant.setup import async_setup_component
+import homeassistant.components.automation as automation
+
+from tests.common import mock_service, mock_coro
+
+
+@asyncio.coroutine
+def test_if_fires_on_hass_start(hass):
+    """Test the firing when HASS starts."""
+    calls = mock_service(hass, 'test', 'automation')
+    hass.state = CoreState.not_running
+    config = {
+        automation.DOMAIN: {
+            'alias': 'hello',
+            'trigger': {
+                'platform': 'homeassistant',
+                'event': 'start',
+            },
+            'action': {
+                'service': 'test.automation',
+            }
+        }
+    }
+
+    res = yield from async_setup_component(hass, automation.DOMAIN, config)
+    assert res
+    assert not automation.is_on(hass, 'automation.hello')
+    assert len(calls) == 0
+
+    yield from hass.async_start()
+    assert automation.is_on(hass, 'automation.hello')
+    assert len(calls) == 1
+
+    with patch('homeassistant.config.async_hass_config_yaml',
+               Mock(return_value=mock_coro(config))):
+        yield from hass.services.async_call(
+            automation.DOMAIN, automation.SERVICE_RELOAD, blocking=True)
+
+    assert automation.is_on(hass, 'automation.hello')
+    assert len(calls) == 1
+
+
+@asyncio.coroutine
+def test_if_fires_on_hass_shutdown(hass):
+    """Test the firing when HASS starts."""
+    calls = mock_service(hass, 'test', 'automation')
+    hass.state = CoreState.not_running
+
+    res = yield from async_setup_component(hass, automation.DOMAIN, {
+        automation.DOMAIN: {
+            'alias': 'hello',
+            'trigger': {
+                'platform': 'homeassistant',
+                'event': 'shutdown',
+            },
+            'action': {
+                'service': 'test.automation',
+            }
+        }
+    })
+    assert res
+    assert not automation.is_on(hass, 'automation.hello')
+    assert len(calls) == 0
+
+    yield from hass.async_start()
+    assert automation.is_on(hass, 'automation.hello')
+    assert len(calls) == 0
+
+    with patch.object(hass.loop, 'stop'):
+        yield from hass.async_stop()
+    assert len(calls) == 1
+
+    # with patch('homeassistant.config.async_hass_config_yaml',
+    #            Mock(return_value=mock_coro(config))):
+    #     yield from hass.services.async_call(
+    #         automation.DOMAIN, automation.SERVICE_RELOAD, blocking=True)
+
+    # assert automation.is_on(hass, 'automation.hello')
+    # assert len(calls) == 1

tests/components/binary_sensor/test_ffmpeg.py

@@ -65,7 +65,7 @@ class TestFFmpegNoiseSetup(object):
         entity = self.hass.states.get('binary_sensor.ffmpeg_noise')
         assert entity.state == 'off'
 
-        mock_ffmpeg.call_args[0][2](True)
+        self.hass.add_job(mock_ffmpeg.call_args[0][2], True)
         self.hass.block_till_done()
 
         entity = self.hass.states.get('binary_sensor.ffmpeg_noise')
@@ -130,7 +130,7 @@ class TestFFmpegMotionSetup(object):
         entity = self.hass.states.get('binary_sensor.ffmpeg_motion')
         assert entity.state == 'off'
 
-        mock_ffmpeg.call_args[0][2](True)
+        self.hass.add_job(mock_ffmpeg.call_args[0][2], True)
         self.hass.block_till_done()
 
         entity = self.hass.states.get('binary_sensor.ffmpeg_motion')

tests/components/test_alert.py

@@ -166,7 +166,7 @@ class TestAlert(unittest.TestCase):
     def test_noack(self):
         """Test no ack feature."""
         entity = alert.Alert(self.hass, *TEST_NOACK)
-        self.hass.async_add_job(entity.begin_alerting)
+        self.hass.add_job(entity.begin_alerting)
         self.hass.block_till_done()
 
         self.assertEqual(True, entity.hidden)

tests/components/zwave/test_node_entity.py

@@ -1,49 +1,33 @@
 """Test Z-Wave node entity."""
+import asyncio
 import unittest
-from unittest.mock import patch, Mock
-from tests.common import get_test_home_assistant
+from unittest.mock import patch
 import tests.mock.zwave as mock_zwave
 import pytest
 from homeassistant.components.zwave import node_entity
 
 
-@pytest.mark.usefixtures('mock_openzwave')
-class TestZWaveBaseEntity(unittest.TestCase):
-    """Class to test ZWaveBaseEntity."""
+@asyncio.coroutine
+def test_maybe_schedule_update(hass, mock_openzwave):
+    """Test maybe schedule update."""
+    base_entity = node_entity.ZWaveBaseEntity()
+    base_entity.hass = hass
 
-    def setUp(self):
-        """Initialize values for this testcase class."""
-        self.hass = get_test_home_assistant()
+    with patch.object(hass.loop, 'call_later') as mock_call_later:
+        base_entity._schedule_update()
+        assert mock_call_later.called
 
-        def call_soon(time, func, *args):
-            """Replace call_later by call_soon."""
-            return self.hass.loop.call_soon(func, *args)
+        base_entity._schedule_update()
+        assert len(mock_call_later.mock_calls) == 1
 
-        self.hass.loop.call_later = call_soon
-        self.base_entity = node_entity.ZWaveBaseEntity()
-        self.base_entity.hass = self.hass
-        self.hass.start()
+        do_update = mock_call_later.mock_calls[0][1][1]
 
-    def tearDown(self):  # pylint: disable=invalid-name
-        """Stop everything that was started."""
-        self.hass.stop()
+        with patch.object(hass, 'async_add_job') as mock_add_job:
+            do_update()
+            assert mock_add_job.called
 
-    def test_maybe_schedule_update(self):
-        """Test maybe_schedule_update."""
-        with patch.object(self.base_entity, 'async_update_ha_state',
-                          Mock()) as mock_update:
-            self.base_entity.maybe_schedule_update()
-            self.hass.block_till_done()
-            mock_update.assert_called_once_with()
-
-    def test_maybe_schedule_update_called_twice(self):
-        """Test maybe_schedule_update called twice."""
-        with patch.object(self.base_entity, 'async_update_ha_state',
-                          Mock()) as mock_update:
-            self.base_entity.maybe_schedule_update()
-            self.base_entity.maybe_schedule_update()
-            self.hass.block_till_done()
-            mock_update.assert_called_once_with()
+        base_entity._schedule_update()
+        assert len(mock_call_later.mock_calls) == 2
 
 
 @pytest.mark.usefixtures('mock_openzwave')