Improve availability of Tractive entities (#97091)

Co-authored-by: Robert Resch <robert@resch.dev>

homeassistant/components/tractive/__init__.py

@@ -89,15 +89,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise ConfigEntryNotReady from error
 
     tractive = TractiveClient(hass, client, creds["user_id"], entry)
-    tractive.subscribe()
 
     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`.
@@ -173,6 +171,14 @@ def user_id(self) -> str:
         """Return user id."""
         return self._user_id
 
+    @property
+    def subscribed(self) -> bool:
+        """Return True if subscribed."""
+        if self._listen_task is None:
+            return False
+
+        return not self._listen_task.cancelled()
+
     async def trackable_objects(
         self,
     ) -> list[aiotractive.trackable_object.TrackableObject]:

homeassistant/components/tractive/binary_sensor.py

@@ -11,63 +11,40 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_BATTERY_CHARGING, EntityCategory
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import Trackables
-from .const import (
-    CLIENT,
-    DOMAIN,
-    SERVER_UNAVAILABLE,
-    TRACKABLES,
-    TRACKER_HARDWARE_STATUS_UPDATED,
-)
+from . import Trackables, TractiveClient
+from .const import CLIENT, DOMAIN, TRACKABLES, TRACKER_HARDWARE_STATUS_UPDATED
 from .entity import TractiveEntity
 
 
 class TractiveBinarySensor(TractiveEntity, BinarySensorEntity):
     """Tractive sensor."""
 
     def __init__(
-        self, user_id: str, item: Trackables, description: BinarySensorEntityDescription
+        self,
+        client: TractiveClient,
+        item: Trackables,
+        description: BinarySensorEntityDescription,
     ) -> None:
         """Initialize sensor entity."""
-        super().__init__(user_id, item.trackable, item.tracker_details)
+        super().__init__(
+            client,
+            item.trackable,
+            item.tracker_details,
+            f"{TRACKER_HARDWARE_STATUS_UPDATED}-{item.tracker_details['_id']}",
+        )
 
         self._attr_unique_id = f"{item.trackable['_id']}_{description.key}"
-        self.entity_description = description
-
-    @callback
-    def handle_server_unavailable(self) -> None:
-        """Handle server unavailable."""
         self._attr_available = False
-        self.async_write_ha_state()
+        self.entity_description = description
 
     @callback
-    def handle_hardware_status_update(self, event: dict[str, Any]) -> None:
-        """Handle hardware status update."""
+    def handle_status_update(self, event: dict[str, Any]) -> None:
+        """Handle status update."""
         self._attr_is_on = event[self.entity_description.key]
-        self._attr_available = True
-        self.async_write_ha_state()
-
-    async def async_added_to_hass(self) -> None:
-        """Handle entity which will be added."""
 
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass,
-                f"{TRACKER_HARDWARE_STATUS_UPDATED}-{self._tracker_id}",
-                self.handle_hardware_status_update,
-            )
-        )
-
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass,
-                f"{SERVER_UNAVAILABLE}-{self._user_id}",
-                self.handle_server_unavailable,
-            )
-        )
+        super().handle_status_update(event)
 
 
 SENSOR_TYPE = BinarySensorEntityDescription(
@@ -86,7 +63,7 @@ async def async_setup_entry(
     trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
 
     entities = [
-        TractiveBinarySensor(client.user_id, item, SENSOR_TYPE)
+        TractiveBinarySensor(client, item, SENSOR_TYPE)
         for item in trackables
         if item.tracker_details.get("charging_state") is not None
     ]

homeassistant/components/tractive/device_tracker.py

@@ -9,7 +9,7 @@
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import Trackables
+from . import Trackables, TractiveClient
 from .const import (
     CLIENT,
     DOMAIN,
@@ -28,7 +28,7 @@ async def async_setup_entry(
     client = hass.data[DOMAIN][entry.entry_id][CLIENT]
     trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
 
-    entities = [TractiveDeviceTracker(client.user_id, item) for item in trackables]
+    entities = [TractiveDeviceTracker(client, item) for item in trackables]
 
     async_add_entities(entities)
 
@@ -39,9 +39,14 @@ class TractiveDeviceTracker(TractiveEntity, TrackerEntity):
     _attr_icon = "mdi:paw"
     _attr_translation_key = "tracker"
 
-    def __init__(self, user_id: str, item: Trackables) -> None:
+    def __init__(self, client: TractiveClient, item: Trackables) -> None:
         """Initialize tracker entity."""
-        super().__init__(user_id, item.trackable, item.tracker_details)
+        super().__init__(
+            client,
+            item.trackable,
+            item.tracker_details,
+            f"{TRACKER_HARDWARE_STATUS_UPDATED}-{item.tracker_details['_id']}",
+        )
 
         self._battery_level: int | None = item.hw_info.get("battery_level")
         self._latitude: float = item.pos_report["latlong"][0]
@@ -94,18 +99,15 @@ def _handle_position_update(self, event: dict[str, Any]) -> None:
         self._attr_available = True
         self.async_write_ha_state()
 
-    @callback
-    def _handle_server_unavailable(self) -> None:
-        self._attr_available = False
-        self.async_write_ha_state()
-
     async def async_added_to_hass(self) -> None:
         """Handle entity which will be added."""
+        if not self._client.subscribed:
+            self._client.subscribe()
 
         self.async_on_remove(
             async_dispatcher_connect(
                 self.hass,
-                f"{TRACKER_HARDWARE_STATUS_UPDATED}-{self._tracker_id}",
+                self._dispatcher_signal,
                 self._handle_hardware_status_update,
             )
         )
@@ -122,6 +124,6 @@ async def async_added_to_hass(self) -> None:
             async_dispatcher_connect(
                 self.hass,
                 f"{SERVER_UNAVAILABLE}-{self._user_id}",
-                self._handle_server_unavailable,
+                self.handle_server_unavailable,
             )
         )

homeassistant/components/tractive/entity.py

@@ -3,10 +3,13 @@
 
 from typing import Any
 
+from homeassistant.core import callback
 from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import Entity
 
-from .const import DOMAIN
+from . import TractiveClient
+from .const import DOMAIN, SERVER_UNAVAILABLE
 
 
 class TractiveEntity(Entity):
@@ -15,7 +18,11 @@ class TractiveEntity(Entity):
     _attr_has_entity_name = True
 
     def __init__(
-        self, user_id: str, trackable: dict[str, Any], tracker_details: dict[str, Any]
+        self,
+        client: TractiveClient,
+        trackable: dict[str, Any],
+        tracker_details: dict[str, Any],
+        dispatcher_signal: str,
     ) -> None:
         """Initialize tracker entity."""
         self._attr_device_info = DeviceInfo(
@@ -26,6 +33,40 @@ def __init__(
             sw_version=tracker_details["fw_version"],
             model=tracker_details["model_number"],
         )
-        self._user_id = user_id
+        self._user_id = client.user_id
         self._tracker_id = tracker_details["_id"]
-        self._trackable = trackable
+        self._client = client
+        self._dispatcher_signal = dispatcher_signal
+
+    async def async_added_to_hass(self) -> None:
+        """Handle entity which will be added."""
+        if not self._client.subscribed:
+            self._client.subscribe()
+
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                self._dispatcher_signal,
+                self.handle_status_update,
+            )
+        )
+
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                f"{SERVER_UNAVAILABLE}-{self._user_id}",
+                self.handle_server_unavailable,
+            )
+        )
+
+    @callback
+    def handle_status_update(self, event: dict[str, Any]) -> None:
+        """Handle status update."""
+        self._attr_available = event[self.entity_description.key] is not None
+        self.async_write_ha_state()
+
+    @callback
+    def handle_server_unavailable(self) -> None:
+        """Handle server unavailable."""
+        self._attr_available = False
+        self.async_write_ha_state()