Move mqtt debouncer to mqtt utils (#120392)

tests/components/mqtt/test_util.py

@@ -1,22 +1,106 @@
 """Test MQTT utils."""
 
+import asyncio
 from collections.abc import Callable
+from datetime import timedelta
 from pathlib import Path
 from random import getrandbits
 import shutil
 import tempfile
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 
 from homeassistant.components import mqtt
+from homeassistant.components.mqtt.models import MessageCallbackType
+from homeassistant.components.mqtt.util import EnsureJobAfterCooldown
 from homeassistant.config_entries import ConfigEntryDisabler, ConfigEntryState
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import CoreState, HomeAssistant
+from homeassistant.util.dt import utcnow
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 from tests.typing import MqttMockHAClient, MqttMockPahoClient
 
 
+async def test_canceling_debouncer_on_shutdown(
+    hass: HomeAssistant,
+    record_calls: MessageCallbackType,
+    setup_with_birth_msg_client_mock: MqttMockPahoClient,
+) -> None:
+    """Test canceling the debouncer when HA shuts down."""
+    mqtt_client_mock = setup_with_birth_msg_client_mock
+
+    with patch("homeassistant.components.mqtt.client.SUBSCRIBE_COOLDOWN", 2):
+        await hass.async_block_till_done()
+        async_fire_time_changed(hass, utcnow() + timedelta(seconds=5))
+        await hass.async_block_till_done()
+        await mqtt.async_subscribe(hass, "test/state1", record_calls)
+        async_fire_time_changed(hass, utcnow() + timedelta(seconds=0.1))
+        # Stop HA so the scheduled debouncer task will be canceled
+        mqtt_client_mock.subscribe.reset_mock()
+        hass.bus.fire(EVENT_HOMEASSISTANT_STOP)
+        await mqtt.async_subscribe(hass, "test/state2", record_calls)
+        async_fire_time_changed(hass, utcnow() + timedelta(seconds=0.1))
+        await mqtt.async_subscribe(hass, "test/state3", record_calls)
+        async_fire_time_changed(hass, utcnow() + timedelta(seconds=0.1))
+        await mqtt.async_subscribe(hass, "test/state4", record_calls)
+        async_fire_time_changed(hass, utcnow() + timedelta(seconds=0.1))
+        await mqtt.async_subscribe(hass, "test/state5", record_calls)
+        async_fire_time_changed(hass, utcnow() + timedelta(seconds=0.1))
+        await hass.async_block_till_done()
+
+        mqtt_client_mock.subscribe.assert_not_called()
+
+        # Note thet the broker connection will not be disconnected gracefully
+        await hass.async_block_till_done()
+        async_fire_time_changed(hass, utcnow() + timedelta(seconds=5))
+        await asyncio.sleep(0)
+        await hass.async_block_till_done(wait_background_tasks=True)
+        mqtt_client_mock.subscribe.assert_not_called()
+        mqtt_client_mock.disconnect.assert_not_called()
+
+
+async def test_canceling_debouncer_normal(
+    hass: HomeAssistant,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test canceling the debouncer before completion."""
+
+    async def _async_myjob() -> None:
+        await asyncio.sleep(1.0)
+
+    debouncer = EnsureJobAfterCooldown(0.0, _async_myjob)
+    debouncer.async_schedule()
+    await asyncio.sleep(0.01)
+    assert debouncer._task is not None
+    await debouncer.async_cleanup()
+    assert debouncer._task is None
+
+
+async def test_canceling_debouncer_throws(
+    hass: HomeAssistant,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test canceling the debouncer when HA shuts down."""
+
+    async def _async_myjob() -> None:
+        await asyncio.sleep(1.0)
+
+    debouncer = EnsureJobAfterCooldown(0.0, _async_myjob)
+    debouncer.async_schedule()
+    await asyncio.sleep(0.01)
+    assert debouncer._task is not None
+    # let debouncer._task fail by mocking it
+    with patch.object(debouncer, "_task") as task:
+        task.cancel = MagicMock(return_value=True)
+        await debouncer.async_cleanup()
+        assert "Error cleaning up task" in caplog.text
+        await hass.async_block_till_done()
+        async_fire_time_changed(hass, utcnow() + timedelta(seconds=5))
+        await hass.async_block_till_done()
+
+
 async def help_create_test_certificate_file(
     hass: HomeAssistant,
     mock_temp_dir: str,