diff --git a/homeassistant/components/imap/__init__.py b/homeassistant/components/imap/__init__.py index 7e582aa04d4..468181be5f7 100644 --- a/homeassistant/components/imap/__init__.py +++ b/homeassistant/components/imap/__init__.py @@ -15,7 +15,11 @@ from homeassistant.exceptions import ( ) from .const import DOMAIN -from .coordinator import ImapDataUpdateCoordinator, connect_to_server +from .coordinator import ( + ImapPollingDataUpdateCoordinator, + ImapPushDataUpdateCoordinator, + connect_to_server, +) from .errors import InvalidAuth, InvalidFolder PLATFORMS: list[Platform] = [Platform.SENSOR] @@ -32,7 +36,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: except (asyncio.TimeoutError, AioImapException) as err: raise ConfigEntryNotReady from err - coordinator = ImapDataUpdateCoordinator(hass, imap_client) + coordinator_class: type[ + ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator + ] + if imap_client.has_capability("IDLE"): + coordinator_class = ImapPushDataUpdateCoordinator + else: + coordinator_class = ImapPollingDataUpdateCoordinator + + coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator = ( + coordinator_class(hass, imap_client) + ) await coordinator.async_config_entry_first_refresh() hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator @@ -49,6 +63,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - coordinator: ImapDataUpdateCoordinator = hass.data[DOMAIN].pop(entry.entry_id) + coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator = hass.data[ + DOMAIN + ].pop( + entry.entry_id + ) await coordinator.shutdown() return unload_ok diff --git a/homeassistant/components/imap/coordinator.py b/homeassistant/components/imap/coordinator.py index e170f79e7f4..e9bbb623013 100644 --- a/homeassistant/components/imap/coordinator.py +++ b/homeassistant/components/imap/coordinator.py @@ -10,9 +10,10 @@ from typing import Any from aioimaplib import AUTH, IMAP4_SSL, SELECTED, AioImapException import async_timeout -from homeassistant.config_entries import ConfigEntry +from homeassistant.config_entries import ConfigEntry, ConfigEntryState from homeassistant.const import CONF_PASSWORD, CONF_PORT, CONF_USERNAME from homeassistant.core import HomeAssistant +from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryError from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import CONF_CHARSET, CONF_FOLDER, CONF_SEARCH, CONF_SERVER, DOMAIN @@ -20,6 +21,8 @@ from .errors import InvalidAuth, InvalidFolder _LOGGER = logging.getLogger(__name__) +BACKOFF_TIME = 10 + async def connect_to_server(data: Mapping[str, Any]) -> IMAP4_SSL: """Connect to imap server and return client.""" @@ -27,82 +30,179 @@ async def connect_to_server(data: Mapping[str, Any]) -> IMAP4_SSL: await client.wait_hello_from_server() await client.login(data[CONF_USERNAME], data[CONF_PASSWORD]) if client.protocol.state != AUTH: - raise InvalidAuth + raise InvalidAuth("Invalid username or password") await client.select(data[CONF_FOLDER]) if client.protocol.state != SELECTED: - raise InvalidFolder + raise InvalidFolder(f"Folder {data[CONF_FOLDER]} is invalid") return client -class ImapDataUpdateCoordinator(DataUpdateCoordinator[int]): - """Class for imap client.""" +class ImapDataUpdateCoordinator(DataUpdateCoordinator[int | None]): + """Base class for imap client.""" config_entry: ConfigEntry - def __init__(self, hass: HomeAssistant, imap_client: IMAP4_SSL) -> None: + def __init__( + self, + hass: HomeAssistant, + imap_client: IMAP4_SSL, + update_interval: timedelta | None, + ) -> None: """Initiate imap client.""" - self.hass = hass self.imap_client = imap_client - self.support_push = imap_client.has_capability("IDLE") super().__init__( hass, _LOGGER, name=DOMAIN, - update_interval=timedelta(seconds=10) if not self.support_push else None, + update_interval=update_interval, ) - async def _async_update_data(self) -> int: - """Update the number of unread emails.""" - try: - if self.imap_client is None: - self.imap_client = await connect_to_server(self.config_entry.data) - except (AioImapException, asyncio.TimeoutError) as err: - raise UpdateFailed(err) from err + async def async_start(self) -> None: + """Start coordinator.""" - return await self.refresh_email_count() - - async def refresh_email_count(self) -> int: - """Check the number of found emails.""" - try: - await self.imap_client.noop() - result, lines = await self.imap_client.search( - self.config_entry.data[CONF_SEARCH], - charset=self.config_entry.data[CONF_CHARSET], - ) - except (AioImapException, asyncio.TimeoutError) as err: - raise UpdateFailed(err) from err + async def _async_reconnect_if_needed(self) -> None: + """Connect to imap server.""" + if self.imap_client is None: + self.imap_client = await connect_to_server(self.config_entry.data) + async def _async_fetch_number_of_messages(self) -> int | None: + """Fetch number of messages.""" + await self._async_reconnect_if_needed() + await self.imap_client.noop() + result, lines = await self.imap_client.search( + self.config_entry.data[CONF_SEARCH], + charset=self.config_entry.data[CONF_CHARSET], + ) if result != "OK": raise UpdateFailed( f"Invalid response for search '{self.config_entry.data[CONF_SEARCH]}': {result} / {lines[0]}" ) - if self.support_push: - self.hass.async_create_background_task( - self.async_wait_server_push(), "Wait for IMAP data push" - ) return len(lines[0].split()) - async def async_wait_server_push(self) -> None: - """Wait for data push from server.""" - try: - idle: asyncio.Future = await self.imap_client.idle_start() - await self.imap_client.wait_server_push() - self.imap_client.idle_done() - async with async_timeout.timeout(10): - await idle - - except (AioImapException, asyncio.TimeoutError): - _LOGGER.warning( - "Lost %s (will attempt to reconnect)", - self.config_entry.data[CONF_SERVER], - ) + async def _cleanup(self, log_error: bool = False) -> None: + """Close resources.""" + if self.imap_client: + try: + if self.imap_client.has_pending_idle(): + self.imap_client.idle_done() + await self.imap_client.stop_wait_server_push() + await self.imap_client.close() + await self.imap_client.logout() + except (AioImapException, asyncio.TimeoutError) as ex: + if log_error: + self.async_set_update_error(ex) + _LOGGER.warning("Error while cleaning up imap connection") self.imap_client = None - await self.async_request_refresh() async def shutdown(self, *_) -> None: """Close resources.""" - if self.imap_client: - if self.imap_client.has_pending_idle(): + await self._cleanup(log_error=True) + + +class ImapPollingDataUpdateCoordinator(ImapDataUpdateCoordinator): + """Class for imap client.""" + + def __init__(self, hass: HomeAssistant, imap_client: IMAP4_SSL) -> None: + """Initiate imap client.""" + super().__init__(hass, imap_client, timedelta(seconds=10)) + + async def _async_update_data(self) -> int | None: + """Update the number of unread emails.""" + try: + return await self._async_fetch_number_of_messages() + except ( + AioImapException, + UpdateFailed, + asyncio.TimeoutError, + ) as ex: + self.async_set_update_error(ex) + await self._cleanup() + raise UpdateFailed() from ex + except InvalidFolder as ex: + _LOGGER.warning("Selected mailbox folder is invalid") + self.async_set_update_error(ex) + await self._cleanup() + raise ConfigEntryError("Selected mailbox folder is invalid.") from ex + except InvalidAuth as ex: + _LOGGER.warning("Username or password incorrect, starting reauthentication") + self.async_set_update_error(ex) + await self._cleanup() + raise ConfigEntryAuthFailed() from ex + + +class ImapPushDataUpdateCoordinator(ImapDataUpdateCoordinator): + """Class for imap client.""" + + def __init__(self, hass: HomeAssistant, imap_client: IMAP4_SSL) -> None: + """Initiate imap client.""" + super().__init__(hass, imap_client, None) + self._push_wait_task: asyncio.Task[None] | None = None + + async def _async_update_data(self) -> int | None: + """Update the number of unread emails.""" + await self.async_start() + return None + + async def async_start(self) -> None: + """Start coordinator.""" + self._push_wait_task = self.hass.async_create_background_task( + self._async_wait_push_loop(), "Wait for IMAP data push" + ) + + async def _async_wait_push_loop(self) -> None: + """Wait for data push from server.""" + while True: + try: + number_of_messages = await self._async_fetch_number_of_messages() + except InvalidAuth as ex: + _LOGGER.warning( + "Username or password incorrect, starting reauthentication" + ) + self.config_entry.async_start_reauth(self.hass) + self.async_set_update_error(ex) + await self._cleanup() + await asyncio.sleep(BACKOFF_TIME) + except InvalidFolder as ex: + _LOGGER.warning("Selected mailbox folder is invalid") + self.config_entry.async_set_state( + self.hass, + ConfigEntryState.SETUP_ERROR, + "Selected mailbox folder is invalid.", + ) + self.async_set_update_error(ex) + await self._cleanup() + await asyncio.sleep(BACKOFF_TIME) + except ( + UpdateFailed, + AioImapException, + asyncio.TimeoutError, + ) as ex: + self.async_set_update_error(ex) + await self._cleanup() + await asyncio.sleep(BACKOFF_TIME) + continue + else: + self.async_set_updated_data(number_of_messages) + try: + idle: asyncio.Future = await self.imap_client.idle_start() + await self.imap_client.wait_server_push() self.imap_client.idle_done() - await self.imap_client.stop_wait_server_push() - await self.imap_client.logout() + async with async_timeout.timeout(10): + await idle + + except (AioImapException, asyncio.TimeoutError): + _LOGGER.warning( + "Lost %s (will attempt to reconnect after %s s)", + self.config_entry.data[CONF_SERVER], + BACKOFF_TIME, + ) + self.async_set_update_error(UpdateFailed("Lost connection")) + await self._cleanup() + await asyncio.sleep(BACKOFF_TIME) + continue + + async def shutdown(self, *_) -> None: + """Close resources.""" + if self._push_wait_task: + self._push_wait_task.cancel() + await super().shutdown() diff --git a/homeassistant/components/imap/sensor.py b/homeassistant/components/imap/sensor.py index 20457209e99..0bccce0c98d 100644 --- a/homeassistant/components/imap/sensor.py +++ b/homeassistant/components/imap/sensor.py @@ -15,7 +15,7 @@ from homeassistant.helpers.issue_registry import IssueSeverity, async_create_iss from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import ImapDataUpdateCoordinator +from . import ImapPollingDataUpdateCoordinator, ImapPushDataUpdateCoordinator from .const import ( CONF_CHARSET, CONF_FOLDER, @@ -69,18 +69,26 @@ async def async_setup_entry( ) -> None: """Set up the Imap sensor.""" - coordinator: ImapDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator = ( + hass.data[DOMAIN][entry.entry_id] + ) async_add_entities([ImapSensor(coordinator)]) -class ImapSensor(CoordinatorEntity[ImapDataUpdateCoordinator], SensorEntity): +class ImapSensor( + CoordinatorEntity[ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator], + SensorEntity, +): """Representation of an IMAP sensor.""" _attr_icon = "mdi:email-outline" _attr_has_entity_name = True - def __init__(self, coordinator: ImapDataUpdateCoordinator) -> None: + def __init__( + self, + coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator, + ) -> None: """Initialize the sensor.""" super().__init__(coordinator) # To be removed when YAML import is removed @@ -95,11 +103,6 @@ class ImapSensor(CoordinatorEntity[ImapDataUpdateCoordinator], SensorEntity): ) @property - def native_value(self) -> int: + def native_value(self) -> int | None: """Return the number of emails found.""" return self.coordinator.data - - async def async_update(self) -> None: - """Check for idle state before updating.""" - if not await self.coordinator.imap_client.stop_wait_server_push(): - await super().async_update()