Cancel and await idle_start future if the task was canceled after an IMAP connection was lost (#168662)

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

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:

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(