Use async_add_hass_job for debouncer (#41449)

homeassistant/helpers/debounce.py

@@ -3,7 +3,7 @@ import asyncio
 from logging import Logger
 from typing import Any, Awaitable, Callable, Optional
 
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import HassJob, HomeAssistant, callback
 
 
 class Debouncer:
@@ -26,12 +26,25 @@ class Debouncer:
         """
         self.hass = hass
         self.logger = logger
-        self.function = function
+        self._function = function
         self.cooldown = cooldown
         self.immediate = immediate
         self._timer_task: Optional[asyncio.TimerHandle] = None
         self._execute_at_end_of_timer: bool = False
         self._execute_lock = asyncio.Lock()
+        self._job: Optional[HassJob] = None if function is None else HassJob(function)
+
+    @property
+    def function(self) -> Optional[Callable[..., Awaitable[Any]]]:
+        """Return the function being wrapped by the Debouncer."""
+        return self._function
+
+    @function.setter
+    def function(self, function: Callable[..., Awaitable[Any]]) -> None:
+        """Update the function being wrapped by the Debouncer."""
+        self._function = function
+        if self._job is None or function != self._job.target:
+            self._job = HassJob(function)
 
     async def async_call(self) -> None:
         """Call the function."""
@@ -57,7 +70,7 @@ class Debouncer:
             if self._timer_task:
                 return
 
-            await self.hass.async_add_job(self.function)  # type: ignore
+            await self.hass.async_add_hass_job(self._job)  # type: ignore
 
             self._schedule_timer()
 
@@ -82,7 +95,7 @@ class Debouncer:
                 return  # type: ignore
 
             try:
-                await self.hass.async_add_job(self.function)  # type: ignore
+                await self.hass.async_add_hass_job(self._job)  # type: ignore
             except Exception:  # pylint: disable=broad-except
                 self.logger.exception("Unexpected exception from %s", self.function)
 

tests/helpers/test_debounce.py

@@ -20,17 +20,22 @@ async def test_immediate_works(hass):
     assert len(calls) == 1
     assert debouncer._timer_task is not None
     assert debouncer._execute_at_end_of_timer is False
+    assert debouncer._job.target == debouncer.function
 
     # Call when cooldown active setting execute at end to True
     await debouncer.async_call()
     assert len(calls) == 1
     assert debouncer._timer_task is not None
     assert debouncer._execute_at_end_of_timer is True
+    assert debouncer._job.target == debouncer.function
 
     # Canceling debounce in cooldown
     debouncer.async_cancel()
     assert debouncer._timer_task is None
     assert debouncer._execute_at_end_of_timer is False
+    assert debouncer._job.target == debouncer.function
+
+    before_job = debouncer._job
 
     # Call and let timer run out
     await debouncer.async_call()
@@ -39,6 +44,8 @@ async def test_immediate_works(hass):
     assert len(calls) == 2
     assert debouncer._timer_task is None
     assert debouncer._execute_at_end_of_timer is False
+    assert debouncer._job.target == debouncer.function
+    assert debouncer._job == before_job
 
     # Test calling doesn't execute/cooldown if currently executing.
     await debouncer._execute_lock.acquire()
@@ -47,6 +54,7 @@ async def test_immediate_works(hass):
     assert debouncer._timer_task is None
     assert debouncer._execute_at_end_of_timer is False
     debouncer._execute_lock.release()
+    assert debouncer._job.target == debouncer.function
 
 
 async def test_not_immediate_works(hass):
@@ -84,6 +92,7 @@ async def test_not_immediate_works(hass):
     assert len(calls) == 1
     assert debouncer._timer_task is not None
     assert debouncer._execute_at_end_of_timer is False
+    assert debouncer._job.target == debouncer.function
 
     # Reset debouncer
     debouncer.async_cancel()
@@ -95,3 +104,65 @@ async def test_not_immediate_works(hass):
     assert debouncer._timer_task is None
     assert debouncer._execute_at_end_of_timer is False
     debouncer._execute_lock.release()
+    assert debouncer._job.target == debouncer.function
+
+
+async def test_immediate_works_with_function_swapped(hass):
+    """Test immediate works and we can change out the function."""
+    calls = []
+
+    one_function = AsyncMock(side_effect=lambda: calls.append(1))
+    two_function = AsyncMock(side_effect=lambda: calls.append(2))
+
+    debouncer = debounce.Debouncer(
+        hass,
+        None,
+        cooldown=0.01,
+        immediate=True,
+        function=one_function,
+    )
+
+    # Call when nothing happening
+    await debouncer.async_call()
+    assert len(calls) == 1
+    assert debouncer._timer_task is not None
+    assert debouncer._execute_at_end_of_timer is False
+    assert debouncer._job.target == debouncer.function
+
+    # Call when cooldown active setting execute at end to True
+    await debouncer.async_call()
+    assert len(calls) == 1
+    assert debouncer._timer_task is not None
+    assert debouncer._execute_at_end_of_timer is True
+    assert debouncer._job.target == debouncer.function
+
+    # Canceling debounce in cooldown
+    debouncer.async_cancel()
+    assert debouncer._timer_task is None
+    assert debouncer._execute_at_end_of_timer is False
+    assert debouncer._job.target == debouncer.function
+
+    before_job = debouncer._job
+    debouncer.function = two_function
+
+    # Call and let timer run out
+    await debouncer.async_call()
+    assert len(calls) == 2
+    assert calls == [1, 2]
+    await debouncer._handle_timer_finish()
+    assert len(calls) == 2
+    assert calls == [1, 2]
+    assert debouncer._timer_task is None
+    assert debouncer._execute_at_end_of_timer is False
+    assert debouncer._job.target == debouncer.function
+    assert debouncer._job != before_job
+
+    # Test calling doesn't execute/cooldown if currently executing.
+    await debouncer._execute_lock.acquire()
+    await debouncer.async_call()
+    assert len(calls) == 2
+    assert calls == [1, 2]
+    assert debouncer._timer_task is None
+    assert debouncer._execute_at_end_of_timer is False
+    debouncer._execute_lock.release()
+    assert debouncer._job.target == debouncer.function