diff --git a/homeassistant/components/imap/coordinator.py b/homeassistant/components/imap/coordinator.py index 157db4da174c..b914388e270e 100644 --- a/homeassistant/components/imap/coordinator.py +++ b/homeassistant/components/imap/coordinator.py @@ -494,6 +494,7 @@ class ImapPushDataUpdateCoordinator(ImapDataUpdateCoordinator): async def _async_wait_push_loop(self) -> None: """Wait for data push from server.""" + idle: asyncio.Future | None = None while True: try: self.number_of_messages = await self._async_fetch_number_of_messages() @@ -527,8 +528,9 @@ class ImapPushDataUpdateCoordinator(ImapDataUpdateCoordinator): else: self.auth_errors = 0 self.async_set_updated_data(self.number_of_messages) + try: - idle: asyncio.Future = await self.imap_client.idle_start() + idle = await self.imap_client.idle_start() await self.imap_client.wait_server_push() self.imap_client.idle_done() async with asyncio.timeout(10): @@ -543,6 +545,24 @@ class ImapPushDataUpdateCoordinator(ImapDataUpdateCoordinator): await self._cleanup() await asyncio.sleep(BACKOFF_TIME) + finally: + # Ensure no pending IDLE future survives + if idle is not None and not idle.done(): + idle.cancel() + _LOGGER.debug( + "Canceling IDLE wait for %s", + self.config_entry.data[CONF_SERVER], + ) + try: + await idle + except asyncio.CancelledError: + if ( + current_task := asyncio.current_task() + ) and current_task.cancelling(): + raise + except AioImapException: + pass + async def shutdown(self, *_: Any) -> None: """Close resources.""" if self._push_wait_task: diff --git a/tests/components/imap/test_init.py b/tests/components/imap/test_init.py index dc5727991c18..55c8499ee5b4 100644 --- a/tests/components/imap/test_init.py +++ b/tests/components/imap/test_init.py @@ -538,6 +538,7 @@ async def test_lost_connection_with_imap_push( ) -> None: """Test error handling when the connection is lost.""" # Mock an error in waiting for a pushed update + mock_imap_protocol.idle_start.return_value = asyncio.Future() mock_imap_protocol.wait_server_push.side_effect = imap_wait_server_push_exception config_entry = MockConfigEntry(domain=DOMAIN, data=MOCK_CONFIG) config_entry.add_to_hass(hass) @@ -550,6 +551,10 @@ async def test_lost_connection_with_imap_push( assert state is not None assert state.state == "0" + async_fire_time_changed(hass, utcnow() + timedelta(seconds=30)) + await hass.async_block_till_done() + assert "Canceling IDLE wait for imap.server.com" in caplog.text + @pytest.mark.parametrize("imap_has_capability", [True], ids=["push"]) async def test_fetch_number_of_messages(