diff --git a/homeassistant/components/cloud/__init__.py b/homeassistant/components/cloud/__init__.py index 8352b566afe..0cb2665767b 100644 --- a/homeassistant/components/cloud/__init__.py +++ b/homeassistant/components/cloud/__init__.py @@ -18,7 +18,7 @@ from homeassistant.const import ( EVENT_HOMEASSISTANT_STOP, Platform, ) -from homeassistant.core import HomeAssistant, ServiceCall, callback +from homeassistant.core import HassJob, HomeAssistant, ServiceCall, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import config_validation as cv, entityfilter from homeassistant.helpers.aiohttp_client import async_get_clientsession @@ -311,7 +311,9 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: async_call_later( hass=hass, delay=timedelta(hours=STARTUP_REPAIR_DELAY), - action=async_startup_repairs, + action=HassJob( + async_startup_repairs, "cloud startup repairs", cancel_on_shutdown=True + ), ) return True diff --git a/homeassistant/core.py b/homeassistant/core.py index 90b44e20693..a47ace48424 100644 --- a/homeassistant/core.py +++ b/homeassistant/core.py @@ -217,13 +217,25 @@ class HassJob(Generic[_P, _R_co]): we run the job. """ - __slots__ = ("job_type", "target", "name") + __slots__ = ("job_type", "target", "name", "_cancel_on_shutdown") - def __init__(self, target: Callable[_P, _R_co], name: str | None = None) -> None: + def __init__( + self, + target: Callable[_P, _R_co], + name: str | None = None, + *, + cancel_on_shutdown: bool | None = None, + ) -> None: """Create a job object.""" self.target = target self.name = name self.job_type = _get_hassjob_callable_job_type(target) + self._cancel_on_shutdown = cancel_on_shutdown + + @property + def cancel_on_shutdown(self) -> bool | None: + """Return if the job should be cancelled on shutdown.""" + return self._cancel_on_shutdown def __repr__(self) -> str: """Return the job.""" @@ -730,6 +742,7 @@ class HomeAssistant: self._tasks.add(task) task.add_done_callback(self._tasks.remove) task.cancel() + self._cancel_cancellable_timers() self.exit_code = exit_code @@ -814,6 +827,20 @@ class HomeAssistant: if self._stopped is not None: self._stopped.set() + def _cancel_cancellable_timers(self) -> None: + """Cancel timer handles marked as cancellable.""" + # pylint: disable-next=protected-access + handles: Iterable[asyncio.TimerHandle] = self.loop._scheduled # type: ignore[attr-defined] + for handle in handles: + if ( + not handle.cancelled() + and (args := handle._args) # pylint: disable=protected-access + # pylint: disable-next=unidiomatic-typecheck + and type(job := args[0]) is HassJob + and job.cancel_on_shutdown + ): + handle.cancel() + def _async_log_running_tasks(self, stage: int) -> None: """Log all running tasks.""" for task in self._tasks: diff --git a/tests/helpers/test_config_entry_flow.py b/tests/helpers/test_config_entry_flow.py index 23f6571e8bc..90d8030be79 100644 --- a/tests/helpers/test_config_entry_flow.py +++ b/tests/helpers/test_config_entry_flow.py @@ -355,7 +355,6 @@ async def test_webhook_config_flow_registers_webhook( assert result["data"]["webhook_id"] is not None -@patch("homeassistant.components.cloud.STARTUP_REPAIR_DELAY", 0) async def test_webhook_create_cloudhook( hass: HomeAssistant, webhook_flow_conf: None ) -> None: @@ -411,7 +410,6 @@ async def test_webhook_create_cloudhook( assert result["require_restart"] is False -@patch("homeassistant.components.cloud.STARTUP_REPAIR_DELAY", 0) async def test_webhook_create_cloudhook_aborts_not_connected( hass: HomeAssistant, webhook_flow_conf: None ) -> None: diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py index cd0d7ef069e..6e411e4e5ea 100644 --- a/tests/test_bootstrap.py +++ b/tests/test_bootstrap.py @@ -551,7 +551,6 @@ async def test_setup_hass_takes_longer_than_log_slow_startup( assert "Waiting on integrations to complete setup" in caplog.text -@patch("homeassistant.components.cloud.STARTUP_REPAIR_DELAY", 0) async def test_setup_hass_invalid_yaml( mock_enable_logging: Mock, mock_is_virtual_env: Mock, @@ -607,7 +606,6 @@ async def test_setup_hass_config_dir_nonexistent( ) -@patch("homeassistant.components.cloud.STARTUP_REPAIR_DELAY", 0) async def test_setup_hass_safe_mode( mock_enable_logging: Mock, mock_is_virtual_env: Mock, @@ -642,7 +640,6 @@ async def test_setup_hass_safe_mode( @pytest.mark.parametrize("hass_config", [{"homeassistant": {"non-existing": 1}}]) -@patch("homeassistant.components.cloud.STARTUP_REPAIR_DELAY", 0) async def test_setup_hass_invalid_core_config( mock_hass_config: None, mock_enable_logging: Mock, @@ -681,7 +678,6 @@ async def test_setup_hass_invalid_core_config( } ], ) -@patch("homeassistant.components.cloud.STARTUP_REPAIR_DELAY", 0) async def test_setup_safe_mode_if_no_frontend( mock_hass_config: None, mock_enable_logging: Mock, diff --git a/tests/test_core.py b/tests/test_core.py index 6167bf6a63b..4d7a93c2887 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -33,7 +33,7 @@ from homeassistant.const import ( __version__, ) import homeassistant.core as ha -from homeassistant.core import HomeAssistant, State +from homeassistant.core import HassJob, HomeAssistant, State from homeassistant.exceptions import ( InvalidEntityFormatError, InvalidStateError, @@ -2081,3 +2081,26 @@ async def test_shutdown_does_not_block_on_shielded_tasks( # Cleanup lingering task after test is done sleep_task.cancel() + + +async def test_cancellable_hassjob(hass: HomeAssistant) -> None: + """Simulate a shutdown, ensure cancellable jobs are cancelled.""" + job = MagicMock() + + @ha.callback + def run_job(job: HassJob) -> None: + """Call the action.""" + hass.async_run_hass_job(job) + + timer1 = hass.loop.call_later( + 60, run_job, HassJob(ha.callback(job), cancel_on_shutdown=True) + ) + timer2 = hass.loop.call_later(60, run_job, HassJob(ha.callback(job))) + + await hass.async_stop() + + assert timer1.cancelled() + assert not timer2.cancelled() + + # Cleanup + timer2.cancel()