Refactor async_call_later to improve performance and reduce conversion loss (#87117)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/helpers/event.py

@@ -1349,9 +1349,24 @@ def async_call_later(
     | Callable[[datetime], Coroutine[Any, Any, None] | None],
 ) -> CALLBACK_TYPE:
     """Add a listener that is called in <delay>."""
-    if not isinstance(delay, timedelta):
-        delay = timedelta(seconds=delay)
-    return async_track_point_in_utc_time(hass, action, dt_util.utcnow() + delay)
+    if isinstance(delay, timedelta):
+        delay = delay.total_seconds()
+
+    @callback
+    def run_action(job: HassJob[[datetime], Coroutine[Any, Any, None] | None]) -> None:
+        """Call the action."""
+        hass.async_run_hass_job(job, time_tracker_utcnow())
+
+    job = action if isinstance(action, HassJob) else HassJob(action)
+    cancel_callback = hass.loop.call_later(delay, run_action, job)
+
+    @callback
+    def unsub_call_later_listener() -> None:
+        """Cancel the call_later."""
+        assert cancel_callback is not None
+        cancel_callback.cancel()
+
+    return unsub_call_later_listener
 
 
 call_later = threaded_listener_factory(async_call_later)

tests/components/rfxtrx/conftest.py

@@ -1,9 +1,9 @@
 """Common test tools."""
 from __future__ import annotations
 
-from datetime import timedelta
 from unittest.mock import patch
 
+from freezegun import freeze_time
 import pytest
 
 from homeassistant.components import rfxtrx
@@ -80,13 +80,12 @@ async def rfxtrx_automatic_fixture(hass, rfxtrx):
 async def timestep(hass):
     """Step system time forward."""
 
-    with patch("homeassistant.core.dt_util.utcnow") as mock_utcnow:
-        mock_utcnow.return_value = utcnow()
+    with freeze_time(utcnow()) as frozen_time:
 
         async def delay(seconds):
             """Trigger delay in system."""
-            mock_utcnow.return_value += timedelta(seconds=seconds)
-            async_fire_time_changed(hass, mock_utcnow.return_value)
+            frozen_time.tick(delta=seconds)
+            async_fire_time_changed(hass)
             await hass.async_block_till_done()
 
         yield delay

tests/components/tomorrowio/test_sensor.py

@@ -3,8 +3,8 @@ from __future__ import annotations
 
 from datetime import datetime
 from typing import Any
-from unittest.mock import patch
 
+from freezegun import freeze_time
 import pytest
 
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
@@ -20,7 +20,7 @@ from homeassistant.components.tomorrowio.const import (
     DOMAIN,
 )
 from homeassistant.components.tomorrowio.sensor import TomorrowioSensorEntityDescription
-from homeassistant.config_entries import SOURCE_USER
+from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY, SOURCE_USER
 from homeassistant.const import ATTR_ATTRIBUTION, CONF_NAME
 from homeassistant.core import HomeAssistant, State, callback
 from homeassistant.helpers.entity_registry import async_get
@@ -29,7 +29,7 @@ from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM
 
 from .const import API_V4_ENTRY_DATA
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 CC_SENSOR_ENTITY_ID = "sensor.tomorrow_io_{}"
 
@@ -110,10 +110,9 @@ async def _setup(
     hass: HomeAssistant, sensors: list[str], config: dict[str, Any]
 ) -> State:
     """Set up entry and return entity state."""
-    with patch(
-        "homeassistant.util.dt.utcnow",
-        return_value=datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC),
-    ):
+    with freeze_time(
+        datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC)
+    ) as frozen_time:
         data = _get_config_schema(hass, SOURCE_USER)(config)
         data[CONF_NAME] = DEFAULT_NAME
         config_entry = MockConfigEntry(
@@ -129,6 +128,10 @@ async def _setup(
         for entity_name in sensors:
             _enable_entity(hass, CC_SENSOR_ENTITY_ID.format(entity_name))
         await hass.async_block_till_done()
+        # the enabled entity state will be fired in RELOAD_AFTER_UPDATE_DELAY
+        frozen_time.tick(delta=RELOAD_AFTER_UPDATE_DELAY)
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
         assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == len(sensors)
 
 

tests/components/tomorrowio/test_weather.py

@@ -3,7 +3,8 @@ from __future__ import annotations
 
 from datetime import datetime
 from typing import Any
-from unittest.mock import patch
+
+from freezegun import freeze_time
 
 from homeassistant.components.tomorrowio.config_flow import (
     _get_config_schema,
@@ -41,7 +42,7 @@ from homeassistant.components.weather import (
     ATTR_WEATHER_WIND_SPEED_UNIT,
     DOMAIN as WEATHER_DOMAIN,
 )
-from homeassistant.config_entries import SOURCE_USER
+from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY, SOURCE_USER
 from homeassistant.const import ATTR_ATTRIBUTION, ATTR_FRIENDLY_NAME, CONF_NAME
 from homeassistant.core import HomeAssistant, State, callback
 from homeassistant.helpers.entity_registry import async_get
@@ -49,7 +50,7 @@ from homeassistant.util import dt as dt_util
 
 from .const import API_V4_ENTRY_DATA
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 @callback
@@ -66,10 +67,9 @@ def _enable_entity(hass: HomeAssistant, entity_name: str) -> None:
 
 async def _setup(hass: HomeAssistant, config: dict[str, Any]) -> State:
     """Set up entry and return entity state."""
-    with patch(
-        "homeassistant.util.dt.utcnow",
-        return_value=datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC),
-    ):
+    with freeze_time(
+        datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC)
+    ) as frozen_time:
         data = _get_config_schema(hass, SOURCE_USER)(config)
         data[CONF_NAME] = DEFAULT_NAME
         config_entry = MockConfigEntry(
@@ -85,6 +85,10 @@ async def _setup(hass: HomeAssistant, config: dict[str, Any]) -> State:
         for entity_name in ("hourly", "nowcast"):
             _enable_entity(hass, f"weather.tomorrow_io_{entity_name}")
         await hass.async_block_till_done()
+        # the enabled entity state will be fired in RELOAD_AFTER_UPDATE_DELAY
+        frozen_time.tick(delta=RELOAD_AFTER_UPDATE_DELAY)
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
         assert len(hass.states.async_entity_ids(WEATHER_DOMAIN)) == 3
 
     return hass.states.get("weather.tomorrow_io_daily")

tests/helpers/test_entity_component.py

@@ -186,14 +186,14 @@ async def test_platform_not_ready(hass):
 
     component = EntityComponent(_LOGGER, DOMAIN, hass)
 
-    await component.async_setup({DOMAIN: {"platform": "mod1"}})
-    await hass.async_block_till_done()
-    assert len(platform1_setup.mock_calls) == 1
-    assert "test_domain.mod1" not in hass.config.components
-
     utcnow = dt_util.utcnow()
 
     with freeze_time(utcnow):
+        await component.async_setup({DOMAIN: {"platform": "mod1"}})
+        await hass.async_block_till_done()
+        assert len(platform1_setup.mock_calls) == 1
+        assert "test_domain.mod1" not in hass.config.components
+
         # Should not trigger attempt 2
         async_fire_time_changed(hass, utcnow + timedelta(seconds=29))
         await hass.async_block_till_done()

tests/helpers/test_event.py

@@ -1,11 +1,14 @@
 """Test event helpers."""
 
 import asyncio
+from collections.abc import Callable
+import contextlib
 from datetime import date, datetime, timedelta
 from unittest.mock import patch
 
 from astral import LocationInfo
 import astral.sun
+import async_timeout
 from freezegun import freeze_time
 import jinja2
 import pytest
@@ -43,7 +46,7 @@ from homeassistant.helpers.template import Template, result_as_boolean
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
-from tests.common import async_fire_time_changed
+from tests.common import async_fire_time_changed, async_fire_time_changed_exact
 
 DEFAULT_TIME_ZONE = dt_util.DEFAULT_TIME_ZONE
 
@@ -4047,64 +4050,91 @@ async def test_periodic_task_leaving_dst_2(hass, freezer):
 
 async def test_call_later(hass):
     """Test calling an action later."""
+    future = asyncio.get_running_loop().create_future()
+    delay = 5
+    delay_tolerance = 0.1
+    schedule_utctime = dt_util.utcnow()
 
-    def action():
-        pass
+    @callback
+    def action(__utcnow: datetime):
+        _current_delay = __utcnow.timestamp() - schedule_utctime.timestamp()
+        future.set_result(delay < _current_delay < (delay + delay_tolerance))
 
-    now = datetime(2017, 12, 19, 15, 40, 0, tzinfo=dt_util.UTC)
+    async_call_later(hass, delay, action)
 
-    with patch(
-        "homeassistant.helpers.event.async_track_point_in_utc_time"
-    ) as mock, patch("homeassistant.util.dt.utcnow", return_value=now):
-        async_call_later(hass, 3, action)
+    async_fire_time_changed_exact(hass, dt_util.utcnow() + timedelta(seconds=delay))
 
-    assert len(mock.mock_calls) == 1
-    p_hass, p_action, p_point = mock.mock_calls[0][1]
-    assert p_hass is hass
-    assert p_action is action
-    assert p_point == now + timedelta(seconds=3)
+    async with async_timeout.timeout(delay + delay_tolerance):
+        assert await future, "callback was called but the delay was wrong"
 
 
 async def test_async_call_later(hass):
     """Test calling an action later."""
+    future = asyncio.get_running_loop().create_future()
+    delay = 5
+    delay_tolerance = 0.1
+    schedule_utctime = dt_util.utcnow()
 
-    def action():
-        pass
+    @callback
+    def action(__utcnow: datetime):
+        _current_delay = __utcnow.timestamp() - schedule_utctime.timestamp()
+        future.set_result(delay < _current_delay < (delay + delay_tolerance))
 
-    now = datetime(2017, 12, 19, 15, 40, 0, tzinfo=dt_util.UTC)
+    remove = async_call_later(hass, delay, action)
 
-    with patch(
-        "homeassistant.helpers.event.async_track_point_in_utc_time"
-    ) as mock, patch("homeassistant.util.dt.utcnow", return_value=now):
-        remove = async_call_later(hass, 3, action)
+    async_fire_time_changed_exact(hass, dt_util.utcnow() + timedelta(seconds=delay))
 
-    assert len(mock.mock_calls) == 1
-    p_hass, p_action, p_point = mock.mock_calls[0][1]
-    assert p_hass is hass
-    assert p_action is action
-    assert p_point == now + timedelta(seconds=3)
-    assert remove is mock()
+    async with async_timeout.timeout(delay + delay_tolerance):
+        assert await future, "callback was called but the delay was wrong"
+    assert isinstance(remove, Callable)
+    remove()
 
 
 async def test_async_call_later_timedelta(hass):
     """Test calling an action later with a timedelta."""
+    future = asyncio.get_running_loop().create_future()
+    delay = 5
+    delay_tolerance = 0.1
+    schedule_utctime = dt_util.utcnow()
 
-    def action():
-        pass
+    @callback
+    def action(__utcnow: datetime):
+        _current_delay = __utcnow.timestamp() - schedule_utctime.timestamp()
+        future.set_result(delay < _current_delay < (delay + delay_tolerance))
 
-    now = datetime(2017, 12, 19, 15, 40, 0, tzinfo=dt_util.UTC)
+    remove = async_call_later(hass, timedelta(seconds=delay), action)
 
-    with patch(
-        "homeassistant.helpers.event.async_track_point_in_utc_time"
-    ) as mock, patch("homeassistant.util.dt.utcnow", return_value=now):
-        remove = async_call_later(hass, timedelta(seconds=3), action)
+    async_fire_time_changed_exact(hass, dt_util.utcnow() + timedelta(seconds=delay))
 
-    assert len(mock.mock_calls) == 1
-    p_hass, p_action, p_point = mock.mock_calls[0][1]
-    assert p_hass is hass
-    assert p_action is action
-    assert p_point == now + timedelta(seconds=3)
-    assert remove is mock()
+    async with async_timeout.timeout(delay + delay_tolerance):
+        assert await future, "callback was called but the delay was wrong"
+    assert isinstance(remove, Callable)
+    remove()
+
+
+async def test_async_call_later_cancel(hass):
+    """Test canceling a call_later action."""
+    future = asyncio.get_running_loop().create_future()
+    delay = 0.25
+    delay_tolerance = 0.1
+
+    @callback
+    def action(__now: datetime):
+        future.set_result(False)
+
+    remove = async_call_later(hass, delay, action)
+    # fast forward time a bit..
+    async_fire_time_changed_exact(
+        hass, dt_util.utcnow() + timedelta(seconds=delay - delay_tolerance)
+    )
+    # and remove before firing
+    remove()
+    # fast forward time beyond scheduled
+    async_fire_time_changed_exact(hass, dt_util.utcnow() + timedelta(seconds=delay))
+
+    with contextlib.suppress(asyncio.TimeoutError):
+        async with async_timeout.timeout(delay + delay_tolerance):
+            assert await future, "callback not canceled"
 
 
 async def test_track_state_change_event_chain_multple_entity(hass):