Rework Tractive integration init (#55741)

* Rework integration init

* Suggested chancge

* Use Trackables class

* Use try..except for trackable_objects

* Check that the pet has tracker linked

homeassistant/components/tractive/__init__.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import asyncio
+from dataclasses import dataclass
 import logging
 
 import aiotractive
@@ -21,9 +22,11 @@ from homeassistant.helpers.dispatcher import async_dispatcher_send
 from .const import (
     ATTR_DAILY_GOAL,
     ATTR_MINUTES_ACTIVE,
+    CLIENT,
     DOMAIN,
     RECONNECT_INTERVAL,
     SERVER_UNAVAILABLE,
+    TRACKABLES,
     TRACKER_ACTIVITY_STATUS_UPDATED,
     TRACKER_HARDWARE_STATUS_UPDATED,
     TRACKER_POSITION_UPDATED,
@@ -35,11 +38,21 @@ PLATFORMS = ["device_tracker", "sensor"]
 _LOGGER = logging.getLogger(__name__)
 
 
+@dataclass
+class Trackables:
+    """A class that describes trackables."""
+
+    trackable: dict | None = None
+    tracker_details: dict | None = None
+    hw_info: dict | None = None
+    pos_report: dict | None = None
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up tractive from a config entry."""
     data = entry.data
 
-    hass.data.setdefault(DOMAIN, {})
+    hass.data.setdefault(DOMAIN, {}).setdefault(entry.entry_id, {})
 
     client = aiotractive.Tractive(
         data[CONF_EMAIL], data[CONF_PASSWORD], session=async_get_clientsession(hass)
@@ -56,7 +69,21 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     tractive = TractiveClient(hass, client, creds["user_id"])
     tractive.subscribe()
 
-    hass.data[DOMAIN][entry.entry_id] = tractive
+    try:
+        trackable_objects = await client.trackable_objects()
+        trackables = await asyncio.gather(
+            *(_generate_trackables(client, item) for item in trackable_objects)
+        )
+    except aiotractive.exceptions.TractiveError as error:
+        await tractive.unsubscribe()
+        raise ConfigEntryNotReady from error
+
+    # When the pet defined in Tractive has no tracker linked we get None as `trackable`.
+    # So we have to remove None values from trackables list.
+    trackables = [item for item in trackables if item]
+
+    hass.data[DOMAIN][entry.entry_id][CLIENT] = tractive
+    hass.data[DOMAIN][entry.entry_id][TRACKABLES] = trackables
 
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
@@ -70,12 +97,30 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
+async def _generate_trackables(client, trackable):
+    """Generate trackables."""
+    trackable = await trackable.details()
+
+    # Check that the pet has tracker linked.
+    if not trackable["device_id"]:
+        return
+
+    tracker = client.tracker(trackable["device_id"])
+
+    tracker_details, hw_info, pos_report = await asyncio.gather(
+        tracker.details(), tracker.hw_info(), tracker.pos_report()
+    )
+
+    return Trackables(trackable, tracker_details, hw_info, pos_report)
+
+
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
-        tractive = hass.data[DOMAIN].pop(entry.entry_id)
+        tractive = hass.data[DOMAIN][entry.entry_id].pop(CLIENT)
         await tractive.unsubscribe()
+        hass.data[DOMAIN].pop(entry.entry_id)
     return unload_ok
 
 

homeassistant/components/tractive/const.py

@@ -9,6 +9,9 @@ RECONNECT_INTERVAL = timedelta(seconds=10)
 ATTR_DAILY_GOAL = "daily_goal"
 ATTR_MINUTES_ACTIVE = "minutes_active"
 
+CLIENT = "client"
+TRACKABLES = "trackables"
+
 TRACKER_HARDWARE_STATUS_UPDATED = f"{DOMAIN}_tracker_hardware_status_updated"
 TRACKER_POSITION_UPDATED = f"{DOMAIN}_tracker_position_updated"
 TRACKER_ACTIVITY_STATUS_UPDATED = f"{DOMAIN}_tracker_activity_updated"

homeassistant/components/tractive/device_tracker.py

@@ -1,6 +1,5 @@
 """Support for Tractive device trackers."""
 
-import asyncio
 import logging
 
 from homeassistant.components.device_tracker import SOURCE_TYPE_GPS
@@ -9,8 +8,10 @@ from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
 from .const import (
+    CLIENT,
     DOMAIN,
     SERVER_UNAVAILABLE,
+    TRACKABLES,
     TRACKER_HARDWARE_STATUS_UPDATED,
     TRACKER_POSITION_UPDATED,
 )
@@ -21,31 +22,25 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up Tractive device trackers."""
-    client = hass.data[DOMAIN][entry.entry_id]
+    client = hass.data[DOMAIN][entry.entry_id][CLIENT]
+    trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
 
-    trackables = await client.trackable_objects()
+    entities = []
 
-    entities = await asyncio.gather(
-        *(create_trackable_entity(client, trackable) for trackable in trackables)
-    )
+    for item in trackables:
+        entities.append(
+            TractiveDeviceTracker(
+                client.user_id,
+                item.trackable,
+                item.tracker_details,
+                item.hw_info,
+                item.pos_report,
+            )
+        )
 
     async_add_entities(entities)
 
 
-async def create_trackable_entity(client, trackable):
-    """Create an entity instance."""
-    trackable = await trackable.details()
-    tracker = client.tracker(trackable["device_id"])
-
-    tracker_details, hw_info, pos_report = await asyncio.gather(
-        tracker.details(), tracker.hw_info(), tracker.pos_report()
-    )
-
-    return TractiveDeviceTracker(
-        client.user_id, trackable, tracker_details, hw_info, pos_report
-    )
-
-
 class TractiveDeviceTracker(TractiveEntity, TrackerEntity):
     """Tractive device tracker."""
 

homeassistant/components/tractive/sensor.py

@@ -1,7 +1,6 @@
 """Support for Tractive sensors."""
 from __future__ import annotations
 
-import asyncio
 from dataclasses import dataclass
 
 from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
@@ -17,8 +16,10 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from .const import (
     ATTR_DAILY_GOAL,
     ATTR_MINUTES_ACTIVE,
+    CLIENT,
     DOMAIN,
     SERVER_UNAVAILABLE,
+    TRACKABLES,
     TRACKER_ACTIVITY_STATUS_UPDATED,
     TRACKER_HARDWARE_STATUS_UPDATED,
 )
@@ -137,29 +138,21 @@ SENSOR_TYPES = (
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up Tractive device trackers."""
-    client = hass.data[DOMAIN][entry.entry_id]
-
-    trackables = await client.trackable_objects()
+    client = hass.data[DOMAIN][entry.entry_id][CLIENT]
+    trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
 
     entities = []
 
-    async def _prepare_sensor_entity(item):
-        """Prepare sensor entities."""
-        trackable = await item.details()
-        tracker = client.tracker(trackable["device_id"])
-        tracker_details = await tracker.details()
+    for item in trackables:
         for description in SENSOR_TYPES:
-            unique_id = f"{trackable['_id']}_{description.key}"
             entities.append(
                 description.entity_class(
                     client.user_id,
-                    trackable,
-                    tracker_details,
-                    unique_id,
+                    item.trackable,
+                    item.tracker_details,
+                    f"{item.trackable['_id']}_{description.key}",
                     description,
                 )
             )
 
-    await asyncio.gather(*(_prepare_sensor_entity(item) for item in trackables))
-
     async_add_entities(entities)