Add multi account support

arjenbos · Sep 27, 2024 · 815b2fc · 815b2fc
1 parent e72f901
commit 815b2fc
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 10 deletions.
diff --git a/custom_components/postnl/__init__.py b/custom_components/postnl/__init__.py
@@ -36,7 +36,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> True:
     except requests.exceptions.ConnectionError as exception:
         raise ConfigEntryNotReady("Unable to retrieve oauth data from PostNL") from exception
 
-    hass.data[DOMAIN][entry.entry_id] = auth
+    hass.data[DOMAIN][entry.entry_id] = {
+        'auth': auth
+    }
 
     _LOGGER.debug('Using access token: %s', auth.access_token)
 
@@ -50,6 +52,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> True:
     if "error" in userinfo:
         raise ConfigEntryNotReady("Error in retrieving user information from PostNL.")
 
+    hass.data[DOMAIN][entry.entry_id]['userinfo'] = userinfo
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True

diff --git a/custom_components/postnl/coordinator.py b/custom_components/postnl/coordinator.py
@@ -29,16 +29,16 @@ def __init__(self, hass: HomeAssistant) -> None:
             update_interval=timedelta(seconds=90),
         )
 
-    async def _async_update_data(self) -> dict[str, tuple[Package]]:
+    async def _async_update_data(self) -> dict[str, list[Package]]:
         _LOGGER.debug('Get API data')
         try:
-            auth: AsyncConfigEntryAuth = self.hass.data[DOMAIN][self.config_entry.entry_id]
+            auth: AsyncConfigEntryAuth = self.hass.data[DOMAIN][self.config_entry.entry_id]['auth']
             await auth.check_and_refresh_token()
 
             self.graphq_api = PostNLGraphql(auth.access_token)
             self.jouw_api = PostNLJouwAPI(auth.access_token)
 
-            data: dict[str, tuple[Package]] = {
+            data: dict[str, list[Package]] = {
                 'receiver': [],
                 'sender': []
             }

diff --git a/custom_components/postnl/sensor.py b/custom_components/postnl/sensor.py
@@ -3,9 +3,12 @@
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.entity_registry import async_get as async_get_entity_registry
 
+from . import DOMAIN
 from .coordinator import PostNLCoordinator
 from .structs.package import Package
 
@@ -18,24 +21,45 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_e
     coordinator = PostNLCoordinator(hass)
 
     await coordinator.async_config_entry_first_refresh()
+    userinfo = hass.data[DOMAIN][entry.entry_id]['userinfo']
+
+    # Backwards compatibility support with old sensor names and to support the legacy lovelace.
+    unique_id_prefix = "postnl_"
+    entity_registry = async_get_entity_registry(hass)
+    if entity_registry.async_get_entity_id(
+        domain="sensor",
+        platform="postnl",
+        unique_id="postnl_delivery"
+    ) or entity_registry.async_get_entity_id(
+        domain="sensor",
+        platform="postnl",
+        unique_id="postnl_distribution"
+    ):
+        unique_id_prefix = userinfo.get('account_id') + "_"
 
     async_add_entities([
         PostNLDelivery(
             coordinator=coordinator,
+            postnl_userinfo=userinfo,
+            unique_id= unique_id_prefix + "delivery",
             name="PostNL_delivery"
         ),
         PostNLDelivery(
             coordinator=coordinator,
+            postnl_userinfo=userinfo,
             name="PostNL_distribution",
+            unique_id=unique_id_prefix + "distribution",
             receiver=False
         )
     ])
 
 
 class PostNLDelivery(CoordinatorEntity, Entity):
-    def __init__(self, coordinator, name, receiver: bool = True):
+    def __init__(self, coordinator, postnl_userinfo, unique_id, name, receiver: bool = True):
         """Initialize the PostNL sensor."""
         super().__init__(coordinator, context=name)
+        self.postnl_userinfo = postnl_userinfo
+        self._unique_id = unique_id
         self._name: str = name
         self._attributes: dict[str, list[Package]] = {
             'enroute': [],
@@ -46,13 +70,24 @@ def __init__(self, coordinator, name, receiver: bool = True):
         self.handle_coordinator_data()
 
     @property
-    def name(self) -> str:
-        """Return the name of the sensor."""
-        return self._name
+    def unique_id(self) -> str | None:
+        """Return the unique id of the sensor."""
+        return self._unique_id
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return the device info."""
+        return DeviceInfo(
+            identifiers={
+                (DOMAIN, self.postnl_userinfo.get('account_id'))
+            },
+            name=self.postnl_userinfo.get('email'),
+            manufacturer="PostNL",
+        )
 
     @property
-    def unique_id(self) -> str:
-        """Return a unique ID."""
+    def name(self) -> str:
+        """Return the name of the sensor."""
         return self._name
 
     @property