Fix coordinator_context (> HA 2022.07.0)

custom_components/pdns/__init__.py

@@ -1,55 +1,75 @@
 """Integrate with PowerDNS service."""
+from __future__ import annotations
+
 import logging
 from datetime import timedelta
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .pdns import PDNS, CannotConnect, TimeoutExpired, PDNSFailed, DetectionFailed
+from .pdns import PDNS, PDNSFailed
 
 DEFAULT_INTERVAL = 15
 DOMAIN = "pdns"
 CONF_PDNSSRV = "pdns_server"
 CONF_ALIAS = "dns_alias"
 _LOGGER = logging.getLogger(__name__)
+PLATFORMS = ["binary_sensor"]
 
 
-async def async_setup(hass, config):
-    """Load configuration for component."""
-    hass.data.setdefault(DOMAIN, {})
-    return True
-
-
-async def async_setup_entry(hass, config_entry):
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Initialize the component."""
-    alias = config_entry.data.get(CONF_ALIAS)
-    username = config_entry.data.get(CONF_USERNAME)
-    password = config_entry.data.get(CONF_PASSWORD)
-    servername = config_entry.data.get(CONF_PDNSSRV)
-
-    session = async_create_clientsession(hass)
-    client = PDNS(servername, alias, username, password, session)
+    hass.data.setdefault(DOMAIN, {})
 
-    coordinator = DataUpdateCoordinator(
+    coordinator = PDNSDataUpdateCoordinator(
         hass,
-        _LOGGER,
-        name=DOMAIN,
-        update_method=client.async_update,
-        update_interval=timedelta(minutes=DEFAULT_INTERVAL),
+        entry,
+        entry.data.get(CONF_PDNSSRV),
+        entry.data.get(CONF_ALIAS),
+        entry.data.get(CONF_USERNAME),
+        entry.data.get(CONF_PASSWORD),
     )
     await coordinator.async_config_entry_first_refresh()
+    if coordinator.data is None:
+        return False
 
-    hass.data[DOMAIN] = coordinator
-    hass.async_create_task(
-        hass.config_entries.async_forward_entry_setup(config_entry, "binary_sensor")
-    )
+    hass.data[DOMAIN][entry.entry_id] = coordinator
+    hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
     return True
 
 
-async def async_unload_entry(hass, config_entry):
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    await hass.config_entries.async_forward_entry_unload(config_entry, "binary_sensor")
-    return True
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+        hass.data[DOMAIN].pop(entry.entry_id)
+    return unload_ok
+
+
+class PDNSDataUpdateCoordinator(DataUpdateCoordinator):
+    """Define an object to fetch datas."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config_entry: ConfigEntry,
+        servername: str,
+        alias: str,
+        username: str,
+        password: str,
+    ) -> None:
+        """Class to manage fetching data API."""
+        super().__init__(
+            hass, _LOGGER, name=DOMAIN, update_interval=timedelta(DEFAULT_INTERVAL)
+        )
+        session = async_create_clientsession(hass)
+        self.pdns_client = PDNS(servername, alias, username, password, session)
+
+    async def _async_update_data(self) -> dict:
+        try:
+            return await self.pdns_client.async_update()
+        except PDNSFailed as error:
+            raise UpdateFailed(error) from error

custom_components/pdns/binary_sensor.py

@@ -3,27 +3,34 @@
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import DOMAIN
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
     """Defer binary sensor setup to the shared sensor module."""
-    coordinator = hass.data[DOMAIN]
+    coordinator = hass.data[DOMAIN][config_entry.entry_id]
     async_add_entities([DyndnsStatus(coordinator)])
 
 
 class DyndnsStatus(CoordinatorEntity, BinarySensorEntity):
-    """Representation of a VocalMsg sensor."""
+    """Representation of a status sensor."""
 
     _attr_device_class = BinarySensorDeviceClass.PROBLEM
     _attr_name = "Dynamic Update"
     _attr_unique_id = "dynamic_update"
 
     def __init__(self, coordinator):
         """Initialize the sensor."""
-        self.coordinator = coordinator
+        super().__init__(coordinator)
 
     @property
     def is_on(self):

custom_components/pdns/manifest.json

@@ -6,5 +6,5 @@
   "issue_tracker": "https://github.com/Cyr-ius/hass-pdns/issues",
   "codeowners": ["@Cyr-ius"],
   "iot_class": "cloud_polling",
-  "version": "2.1.9"
+  "version": "2.2"
 }