From b9da11030149ace9a341c9e5567c0bd550e4ab43 Mon Sep 17 00:00:00 2001
From: cyr-ius <cyr-ius@ipocus.net>
Date: Sat, 9 Jul 2022 13:58:53 +0200
Subject: [PATCH] Fix coordinator_context (> HA 2022.07.0)

---
 custom_components/pdns/__init__.py      | 74 ++++++++++++++++---------
 custom_components/pdns/binary_sensor.py | 15 +++--
 custom_components/pdns/manifest.json    |  2 +-
 3 files changed, 59 insertions(+), 32 deletions(-)

diff --git a/custom_components/pdns/__init__.py b/custom_components/pdns/__init__.py
index b1f9215..8d57655 100644
--- a/custom_components/pdns/__init__.py
+++ b/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
diff --git a/custom_components/pdns/binary_sensor.py b/custom_components/pdns/binary_sensor.py
index 9f97e52..8437c54 100644
--- a/custom_components/pdns/binary_sensor.py
+++ b/custom_components/pdns/binary_sensor.py
@@ -3,19 +3,26 @@
     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"
@@ -23,7 +30,7 @@ class DyndnsStatus(CoordinatorEntity, BinarySensorEntity):
 
     def __init__(self, coordinator):
         """Initialize the sensor."""
-        self.coordinator = coordinator
+        super().__init__(coordinator)
 
     @property
     def is_on(self):
diff --git a/custom_components/pdns/manifest.json b/custom_components/pdns/manifest.json
index 82369b0..9f24980 100644
--- a/custom_components/pdns/manifest.json
+++ b/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"
 }
\ No newline at end of file