Improve availability of Tractive entities (#97091)

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

homeassistant/components/tractive/__init__.py

@@ -89,7 +89,6 @@ 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()
@@ -97,7 +96,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             *(_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 @@ class TractiveClient:
         """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,17 +11,10 @@ from homeassistant.components.binary_sensor import (
 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
 
 
@@ -29,45 +22,29 @@ 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._attr_available = False
         self.entity_description = description
 
     @callback
-    def handle_server_unavailable(self) -> None:
-        """Handle server unavailable."""
-        self._attr_available = False
-        self.async_write_ha_state()
-
-    @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.core import HomeAssistant, callback
 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 @@ class TractiveDeviceTracker(TractiveEntity, TrackerEntity):
         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 @@ class TractiveDeviceTracker(TractiveEntity, TrackerEntity):
             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 __future__ import annotations
 
 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 @@ class TractiveEntity(Entity):
             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()

homeassistant/components/tractive/sensor.py

@@ -18,10 +18,9 @@ from homeassistant.const import (
     UnitOfTime,
 )
 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 . import Trackables, TractiveClient
 from .const import (
     ATTR_CALORIES,
     ATTR_DAILY_GOAL,
@@ -32,7 +31,6 @@ from .const import (
     ATTR_TRACKER_STATE,
     CLIENT,
     DOMAIN,
-    SERVER_UNAVAILABLE,
     TRACKABLES,
     TRACKER_ACTIVITY_STATUS_UPDATED,
     TRACKER_HARDWARE_STATUS_UPDATED,
@@ -45,7 +43,7 @@ from .entity import TractiveEntity
 class TractiveRequiredKeysMixin:
     """Mixin for required keys."""
 
-    entity_class: type[TractiveSensor]
+    signal_prefix: str
 
 
 @dataclass
@@ -54,112 +52,39 @@ class TractiveSensorEntityDescription(
 ):
     """Class describing Tractive sensor entities."""
 
+    hardware_sensor: bool = False
+
 
 class TractiveSensor(TractiveEntity, SensorEntity):
     """Tractive sensor."""
 
     def __init__(
         self,
-        user_id: str,
+        client: TractiveClient,
         item: Trackables,
         description: TractiveSensorEntityDescription,
     ) -> None:
         """Initialize sensor entity."""
-        super().__init__(user_id, item.trackable, item.tracker_details)
+        if description.hardware_sensor:
+            dispatcher_signal = (
+                f"{description.signal_prefix}-{item.tracker_details['_id']}"
+            )
+        else:
+            dispatcher_signal = f"{description.signal_prefix}-{item.trackable['_id']}"
+        super().__init__(
+            client, item.trackable, item.tracker_details, dispatcher_signal
+        )
 
         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()
-
-
-class TractiveHardwareSensor(TractiveSensor):
-    """Tractive hardware sensor."""
-
-    @callback
-    def handle_hardware_status_update(self, event: dict[str, Any]) -> None:
-        """Handle hardware status update."""
-        if (_state := event[self.entity_description.key]) is None:
-            return
-        self._attr_native_value = _state
-        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,
-            )
-        )
-
-
-class TractiveActivitySensor(TractiveSensor):
-    """Tractive active sensor."""
+        self.entity_description = description
 
     @callback
     def handle_status_update(self, event: dict[str, Any]) -> None:
         """Handle status update."""
         self._attr_native_value = 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_ACTIVITY_STATUS_UPDATED}-{self._trackable['_id']}",
-                self.handle_status_update,
-            )
-        )
-
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass,
-                f"{SERVER_UNAVAILABLE}-{self._user_id}",
-                self.handle_server_unavailable,
-            )
-        )
-
-
-class TractiveWellnessSensor(TractiveActivitySensor):
-    """Tractive wellness sensor."""
-
-    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_WELLNESS_STATUS_UPDATED}-{self._trackable['_id']}",
-                self.handle_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_TYPES: tuple[TractiveSensorEntityDescription, ...] = (
@@ -168,13 +93,15 @@ SENSOR_TYPES: tuple[TractiveSensorEntityDescription, ...] = (
         translation_key="tracker_battery_level",
         native_unit_of_measurement=PERCENTAGE,
         device_class=SensorDeviceClass.BATTERY,
-        entity_class=TractiveHardwareSensor,
+        signal_prefix=TRACKER_HARDWARE_STATUS_UPDATED,
+        hardware_sensor=True,
         entity_category=EntityCategory.DIAGNOSTIC,
     ),
     TractiveSensorEntityDescription(
         key=ATTR_TRACKER_STATE,
         translation_key="tracker_state",
-        entity_class=TractiveHardwareSensor,
+        signal_prefix=TRACKER_HARDWARE_STATUS_UPDATED,
+        hardware_sensor=True,
         icon="mdi:radar",
         entity_category=EntityCategory.DIAGNOSTIC,
         device_class=SensorDeviceClass.ENUM,
@@ -190,7 +117,7 @@ SENSOR_TYPES: tuple[TractiveSensorEntityDescription, ...] = (
         translation_key="activity_time",
         icon="mdi:clock-time-eight-outline",
         native_unit_of_measurement=UnitOfTime.MINUTES,
-        entity_class=TractiveActivitySensor,
+        signal_prefix=TRACKER_ACTIVITY_STATUS_UPDATED,
         state_class=SensorStateClass.TOTAL,
     ),
     TractiveSensorEntityDescription(
@@ -198,7 +125,7 @@ SENSOR_TYPES: tuple[TractiveSensorEntityDescription, ...] = (
         translation_key="rest_time",
         icon="mdi:clock-time-eight-outline",
         native_unit_of_measurement=UnitOfTime.MINUTES,
-        entity_class=TractiveWellnessSensor,
+        signal_prefix=TRACKER_WELLNESS_STATUS_UPDATED,
         state_class=SensorStateClass.TOTAL,
     ),
     TractiveSensorEntityDescription(
@@ -206,7 +133,7 @@ SENSOR_TYPES: tuple[TractiveSensorEntityDescription, ...] = (
         translation_key="calories",
         icon="mdi:fire",
         native_unit_of_measurement="kcal",
-        entity_class=TractiveWellnessSensor,
+        signal_prefix=TRACKER_WELLNESS_STATUS_UPDATED,
         state_class=SensorStateClass.TOTAL,
     ),
     TractiveSensorEntityDescription(
@@ -214,14 +141,14 @@ SENSOR_TYPES: tuple[TractiveSensorEntityDescription, ...] = (
         translation_key="daily_goal",
         icon="mdi:flag-checkered",
         native_unit_of_measurement=UnitOfTime.MINUTES,
-        entity_class=TractiveActivitySensor,
+        signal_prefix=TRACKER_ACTIVITY_STATUS_UPDATED,
     ),
     TractiveSensorEntityDescription(
         key=ATTR_MINUTES_DAY_SLEEP,
         translation_key="minutes_day_sleep",
         icon="mdi:sleep",
         native_unit_of_measurement=UnitOfTime.MINUTES,
-        entity_class=TractiveWellnessSensor,
+        signal_prefix=TRACKER_WELLNESS_STATUS_UPDATED,
         state_class=SensorStateClass.TOTAL,
     ),
     TractiveSensorEntityDescription(
@@ -229,7 +156,7 @@ SENSOR_TYPES: tuple[TractiveSensorEntityDescription, ...] = (
         translation_key="minutes_night_sleep",
         icon="mdi:sleep",
         native_unit_of_measurement=UnitOfTime.MINUTES,
-        entity_class=TractiveWellnessSensor,
+        signal_prefix=TRACKER_WELLNESS_STATUS_UPDATED,
         state_class=SensorStateClass.TOTAL,
     ),
 )
@@ -243,7 +170,7 @@ async def async_setup_entry(
     trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
 
     entities = [
-        description.entity_class(client.user_id, item, description)
+        TractiveSensor(client, item, description)
         for description in SENSOR_TYPES
         for item in trackables
     ]

homeassistant/components/tractive/switch.py

@@ -11,17 +11,15 @@ from homeassistant.components.switch import SwitchEntity, SwitchEntityDescriptio
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import 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 . import Trackables, TractiveClient
 from .const import (
     ATTR_BUZZER,
     ATTR_LED,
     ATTR_LIVE_TRACKING,
     CLIENT,
     DOMAIN,
-    SERVER_UNAVAILABLE,
     TRACKABLES,
     TRACKER_HARDWARE_STATUS_UPDATED,
 )
@@ -77,7 +75,7 @@ async def async_setup_entry(
     trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
 
     entities = [
-        TractiveSwitch(client.user_id, item, description)
+        TractiveSwitch(client, item, description)
         for description in SWITCH_TYPES
         for item in trackables
     ]
@@ -92,12 +90,17 @@ class TractiveSwitch(TractiveEntity, SwitchEntity):
 
     def __init__(
         self,
-        user_id: str,
+        client: TractiveClient,
         item: Trackables,
         description: TractiveSwitchEntityDescription,
     ) -> None:
         """Initialize switch 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._attr_available = False
@@ -106,38 +109,11 @@ class TractiveSwitch(TractiveEntity, SwitchEntity):
         self.entity_description = description
 
     @callback
-    def handle_server_unavailable(self) -> None:
-        """Handle server unavailable."""
-        self._attr_available = False
-        self.async_write_ha_state()
+    def handle_status_update(self, event: dict[str, Any]) -> None:
+        """Handle status update."""
+        self._attr_is_on = event[self.entity_description.key]
 
-    @callback
-    def handle_hardware_status_update(self, event: dict[str, Any]) -> None:
-        """Handle hardware status update."""
-        if (state := event[self.entity_description.key]) is None:
-            return
-        self._attr_is_on = state
-        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)
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn on a switch."""