Improve imap error handling for config entry (#96724)

* Improve error handling config entry

* Removed CancelledError

* Add cleanup

* Do not call protected async_set_state()

homeassistant/components/imap/coordinator.py

@@ -13,7 +13,7 @@ from typing import Any
 from aioimaplib import AUTH, IMAP4_SSL, NONAUTH, SELECTED, AioImapException
 import async_timeout
 
-from homeassistant.config_entries import ConfigEntry, ConfigEntryState
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_PASSWORD,
     CONF_PORT,
@@ -54,6 +54,7 @@ _LOGGER = logging.getLogger(__name__)
 BACKOFF_TIME = 10
 
 EVENT_IMAP = "imap_content"
+MAX_ERRORS = 3
 MAX_EVENT_DATA_BYTES = 32168
 
 
@@ -174,6 +175,7 @@ class ImapDataUpdateCoordinator(DataUpdateCoordinator[int | None]):
     ) -> None:
         """Initiate imap client."""
         self.imap_client = imap_client
+        self.auth_errors: int = 0
         self._last_message_id: str | None = None
         self.custom_event_template = None
         _custom_event_template = entry.data.get(CONF_CUSTOM_EVENT_DATA_TEMPLATE)
@@ -315,7 +317,9 @@ class ImapPollingDataUpdateCoordinator(ImapDataUpdateCoordinator):
     async def _async_update_data(self) -> int | None:
         """Update the number of unread emails."""
         try:
-            return await self._async_fetch_number_of_messages()
+            messages = await self._async_fetch_number_of_messages()
+            self.auth_errors = 0
+            return messages
         except (
             AioImapException,
             UpdateFailed,
@@ -330,8 +334,15 @@ class ImapPollingDataUpdateCoordinator(ImapDataUpdateCoordinator):
             self.async_set_update_error(ex)
             raise ConfigEntryError("Selected mailbox folder is invalid.") from ex
         except InvalidAuth as ex:
-            _LOGGER.warning("Username or password incorrect, starting reauthentication")
             await self._cleanup()
+            self.auth_errors += 1
+            if self.auth_errors <= MAX_ERRORS:
+                _LOGGER.warning("Authentication failed, retrying")
+            else:
+                _LOGGER.warning(
+                    "Username or password incorrect, starting reauthentication"
+                )
+                self.config_entry.async_start_reauth(self.hass)
             self.async_set_update_error(ex)
             raise ConfigEntryAuthFailed() from ex
 
@@ -359,27 +370,28 @@ class ImapPushDataUpdateCoordinator(ImapDataUpdateCoordinator):
 
     async def _async_wait_push_loop(self) -> None:
         """Wait for data push from server."""
+        cleanup = False
         while True:
             try:
                 number_of_messages = await self._async_fetch_number_of_messages()
             except InvalidAuth as ex:
+                self.auth_errors += 1
                 await self._cleanup()
-                _LOGGER.warning(
-                    "Username or password incorrect, starting reauthentication"
-                )
-                self.config_entry.async_start_reauth(self.hass)
+                if self.auth_errors <= MAX_ERRORS:
+                    _LOGGER.warning("Authentication failed, retrying")
+                else:
+                    _LOGGER.warning(
+                        "Username or password incorrect, starting reauthentication"
+                    )
+                    self.config_entry.async_start_reauth(self.hass)
                 self.async_set_update_error(ex)
                 await asyncio.sleep(BACKOFF_TIME)
             except InvalidFolder as ex:
                 _LOGGER.warning("Selected mailbox folder is invalid")
                 await self._cleanup()
-                self.config_entry._async_set_state(  # pylint: disable=protected-access
-                    self.hass,
-                    ConfigEntryState.SETUP_ERROR,
-                    "Selected mailbox folder is invalid.",
-                )
                 self.async_set_update_error(ex)
                 await asyncio.sleep(BACKOFF_TIME)
+                continue
             except (
                 UpdateFailed,
                 AioImapException,
@@ -390,6 +402,7 @@ class ImapPushDataUpdateCoordinator(ImapDataUpdateCoordinator):
                 await asyncio.sleep(BACKOFF_TIME)
                 continue
             else:
+                self.auth_errors = 0
                 self.async_set_updated_data(number_of_messages)
             try:
                 idle: asyncio.Future = await self.imap_client.idle_start()
@@ -398,6 +411,10 @@ class ImapPushDataUpdateCoordinator(ImapDataUpdateCoordinator):
                 async with async_timeout.timeout(10):
                     await idle
 
+            # From python 3.11 asyncio.TimeoutError is an alias of TimeoutError
+            except asyncio.CancelledError as ex:
+                cleanup = True
+                raise asyncio.CancelledError from ex
             except (AioImapException, asyncio.TimeoutError):
                 _LOGGER.debug(
                     "Lost %s (will attempt to reconnect after %s s)",
@@ -406,6 +423,9 @@ class ImapPushDataUpdateCoordinator(ImapDataUpdateCoordinator):
                 )
                 await self._cleanup()
                 await asyncio.sleep(BACKOFF_TIME)
+            finally:
+                if cleanup:
+                    await self._cleanup()
 
     async def shutdown(self, *_: Any) -> None:
         """Close resources."""

tests/components/imap/test_init.py

@@ -235,6 +235,48 @@ async def test_initial_invalid_folder_error(
     assert (state is not None) == success
 
 
+@patch("homeassistant.components.imap.coordinator.MAX_ERRORS", 1)
+@pytest.mark.parametrize("imap_has_capability", [True, False], ids=["push", "poll"])
+async def test_late_authentication_retry(
+    hass: HomeAssistant,
+    caplog: pytest.LogCaptureFixture,
+    mock_imap_protocol: MagicMock,
+) -> None:
+    """Test retrying authentication after a search was failed."""
+
+    # Mock an error in waiting for a pushed update
+    mock_imap_protocol.wait_server_push.side_effect = AioImapException(
+        "Something went wrong"
+    )
+
+    config_entry = MockConfigEntry(domain=DOMAIN, data=MOCK_CONFIG)
+    config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+
+    async_fire_time_changed(hass, utcnow() + timedelta(seconds=60))
+    await hass.async_block_till_done()
+
+    # Mock that the search fails, this will trigger
+    # that the connection will be restarted
+    # Then fail selecting the folder
+    mock_imap_protocol.search.return_value = Response(*BAD_RESPONSE)
+    mock_imap_protocol.login.side_effect = Response(*BAD_RESPONSE)
+
+    async_fire_time_changed(hass, utcnow() + timedelta(seconds=60))
+    await hass.async_block_till_done()
+
+    async_fire_time_changed(hass, utcnow() + timedelta(seconds=60))
+    await hass.async_block_till_done()
+    assert "Authentication failed, retrying" in caplog.text
+
+    # we still should have an entity with an unavailable state
+    state = hass.states.get("sensor.imap_email_email_com")
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE
+
+
+@patch("homeassistant.components.imap.coordinator.MAX_ERRORS", 0)
 @pytest.mark.parametrize("imap_has_capability", [True, False], ids=["push", "poll"])
 async def test_late_authentication_error(
     hass: HomeAssistant,