Rework DataUpdateCoordinator (#404)

custom_components/mail_and_packages/__init__.py

@@ -2,11 +2,11 @@
 import logging
 from datetime import timedelta
 
-import async_timeout
+from async_timeout import timeout
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_RESOURCES
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import (
     CONF_ALLOW_EXTERNAL,
@@ -72,18 +72,13 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     config_entry.options = config_entry.data
     config = config_entry.data
 
-    async def async_update_data():
-        """Fetch data """
-        async with async_timeout.timeout(config.get(CONF_IMAP_TIMEOUT)):
-            return await hass.async_add_executor_job(process_emails, hass, config)
-
-    coordinator = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        name=f"Mail and Packages ({config.get(CONF_HOST)})",
-        update_method=async_update_data,
-        update_interval=timedelta(minutes=config_entry.data.get(CONF_SCAN_INTERVAL)),
-    )
+    # Variables for data coordinator
+    host = config.get(CONF_HOST)
+    timeout = config.get(CONF_IMAP_TIMEOUT)
+    interval = config.get(CONF_SCAN_INTERVAL)
+
+    # Setup the data coordinator
+    coordinator = MailDataUpdateCoordinator(hass, host, timeout, interval, config)
 
     # Fetch initial data so we have data when entities subscribe
     await coordinator.async_refresh()
@@ -187,3 +182,30 @@ async def async_migrate_entry(hass, config_entry):
         _LOGGER.debug("Migration to version %s complete", config_entry.version)
 
     return True
+
+
+class MailDataUpdateCoordinator(DataUpdateCoordinator):
+    """Class to manage fetching mail data."""
+
+    def __init__(self, hass, host, timeout, interval, config):
+        """Initialize."""
+        self.interval = timedelta(minutes=interval)
+        self.name = f"Mail and Packages ({host})"
+        self.timeout = timeout
+        self.config = config
+        self.hass = hass
+
+        _LOGGER.debug("Data will be update every %s", self.interval)
+
+        super().__init__(hass, _LOGGER, name=self.name, update_interval=self.interval)
+
+    async def _async_update_data(self):
+        """Fetch data """
+        async with timeout(self.timeout):
+            try:
+                data = await self.hass.async_add_executor_job(
+                    process_emails, self.hass, self.config
+                )
+            except Exception as error:
+                raise UpdateFailed(error) from error
+            return data

custom_components/mail_and_packages/helpers.py

@@ -12,7 +12,7 @@
 import uuid
 from email.header import decode_header
 from shutil import copyfile, which
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional, Type, Union
 
 import aiohttp
 import imageio as io
@@ -268,7 +268,9 @@ def fetch(hass: Any, config: Any, account: Any, data: dict, sensor: str) -> int:
     return count[sensor]
 
 
-def login(host, port, user, pwd):
+def login(
+    host: str, port: int, user: str, pwd: str
+) -> Union[bool, Type[imaplib.IMAP4_SSL]]:
     """function used to login"""
 
     # Catch invalid mail server / host names
@@ -289,7 +291,7 @@ def login(host, port, user, pwd):
     return account
 
 
-def selectfolder(account, folder) -> None:
+def selectfolder(account: Any, folder: str) -> None:
     """Select folder inside the mailbox"""
     try:
         rv, mailboxes = account.list()

custom_components/mail_and_packages/sensor.py

@@ -99,19 +99,6 @@ def device_state_attributes(self) -> Optional[str]:
             attr[const.ATTR_TRACKING_NUM] = data[tracking]
         return attr
 
-    async def async_update(self):
-        """Update the entity.
-
-        Only used by the generic entity update service.
-        """
-        await self.coordinator.async_request_refresh()
-
-    async def async_added_to_hass(self):
-        """When entity is added to hass."""
-        self.async_on_remove(
-            self.coordinator.async_add_listener(self.async_write_ha_state)
-        )
-
 
 class ImagePathSensors(CoordinatorEntity):
     """ Represntation of a sensor """
@@ -191,16 +178,3 @@ def device_state_attributes(self) -> Optional[str]:
         """Return device specific state attributes."""
         attr = {}
         return attr
-
-    async def async_update(self):
-        """Update the entity.
-
-        Only used by the generic entity update service.
-        """
-        await self.coordinator.async_request_refresh()
-
-    async def async_added_to_hass(self):
-        """When entity is added to hass."""
-        self.async_on_remove(
-            self.coordinator.async_add_listener(self.async_write_ha_state)
-        )

tests/test_helpers.py

@@ -718,10 +718,7 @@ async def test_generate_mp4(
 async def test_connection_error(caplog):
     result = login("localhost", 993, "fakeuser", "suchfakemuchpassword")
     assert not result
-    assert (
-        "Network error while connecting to server: [Errno 111] Connection refused"
-        in caplog.text
-    )
+    assert "Network error while connecting to server:" in caplog.text
 
 
 async def test_login_error(mock_imap_login_error, caplog):