Store Tractive data in config_entry.runtime_data (#116781)

Co-authored-by: J. Nick Koston <nick@koston.org>
Co-authored-by: Maciej Bieniek <478555+bieniu@users.noreply.github.com>
Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>

homeassistant/components/tractive/__init__.py

@@ -33,13 +33,10 @@ from .const import (
     ATTR_MINUTES_REST,
     ATTR_SLEEP_LABEL,
     ATTR_TRACKER_STATE,
-    CLIENT,
     CLIENT_ID,
-    DOMAIN,
     RECONNECT_INTERVAL,
     SERVER_UNAVAILABLE,
     SWITCH_KEY_MAP,
-    TRACKABLES,
     TRACKER_HARDWARE_STATUS_UPDATED,
     TRACKER_POSITION_UPDATED,
     TRACKER_SWITCH_STATUS_UPDATED,
@@ -68,12 +65,21 @@ class Trackables:
     pos_report: dict
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+@dataclass(slots=True)
+class TractiveData:
+    """Class for Tractive data."""
+
+    client: TractiveClient
+    trackables: list[Trackables]
+
+
+TractiveConfigEntry = ConfigEntry[TractiveData]
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: TractiveConfigEntry) -> bool:
     """Set up tractive from a config entry."""
     data = entry.data
 
-    hass.data.setdefault(DOMAIN, {}).setdefault(entry.entry_id, {})
-
     client = aiotractive.Tractive(
         data[CONF_EMAIL],
         data[CONF_PASSWORD],
@@ -101,10 +107,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     # 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]
+    filtered_trackables = [item for item in trackables if item]
 
-    hass.data[DOMAIN][entry.entry_id][CLIENT] = tractive
-    hass.data[DOMAIN][entry.entry_id][TRACKABLES] = trackables
+    entry.runtime_data = TractiveData(tractive, filtered_trackables)
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -114,6 +119,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, cancel_listen_task)
     )
+    entry.async_on_unload(tractive.unsubscribe)
 
     return True
 
@@ -145,14 +151,9 @@ async def _generate_trackables(
     return Trackables(tracker, trackable, tracker_details, hw_info, pos_report)
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: TractiveConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        tractive = hass.data[DOMAIN][entry.entry_id].pop(CLIENT)
-        await tractive.unsubscribe()
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
 class TractiveClient:

homeassistant/components/tractive/binary_sensor.py

@@ -9,13 +9,12 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_BATTERY_CHARGING, EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import Trackables, TractiveClient
-from .const import CLIENT, DOMAIN, TRACKABLES, TRACKER_HARDWARE_STATUS_UPDATED
+from . import Trackables, TractiveClient, TractiveConfigEntry
+from .const import TRACKER_HARDWARE_STATUS_UPDATED
 from .entity import TractiveEntity
 
 
@@ -57,11 +56,13 @@ SENSOR_TYPE = BinarySensorEntityDescription(
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: TractiveConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Tractive device trackers."""
-    client = hass.data[DOMAIN][entry.entry_id][CLIENT]
-    trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
+    client = entry.runtime_data.client
+    trackables = entry.runtime_data.trackables
 
     entities = [
         TractiveBinarySensor(client, item, SENSOR_TYPE)

homeassistant/components/tractive/const.py

@@ -23,9 +23,6 @@ ATTR_TRACKER_STATE = "tracker_state"
 # Please do not use it anywhere else.
 CLIENT_ID = "625e5349c3c3b41c28a669f1"
 
-CLIENT = "client"
-TRACKABLES = "trackables"
-
 TRACKER_HARDWARE_STATUS_UPDATED = f"{DOMAIN}_tracker_hardware_status_updated"
 TRACKER_POSITION_UPDATED = f"{DOMAIN}_tracker_position_updated"
 TRACKER_SWITCH_STATUS_UPDATED = f"{DOMAIN}_tracker_switch_updated"

homeassistant/components/tractive/device_tracker.py

@@ -5,17 +5,13 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.device_tracker import SourceType, TrackerEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import Trackables, TractiveClient
+from . import Trackables, TractiveClient, TractiveConfigEntry
 from .const import (
-    CLIENT,
-    DOMAIN,
     SERVER_UNAVAILABLE,
-    TRACKABLES,
     TRACKER_HARDWARE_STATUS_UPDATED,
     TRACKER_POSITION_UPDATED,
 )
@@ -23,11 +19,13 @@ from .entity import TractiveEntity
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: TractiveConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Tractive device trackers."""
-    client = hass.data[DOMAIN][entry.entry_id][CLIENT]
-    trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
+    client = entry.runtime_data.client
+    trackables = entry.runtime_data.trackables
 
     entities = [TractiveDeviceTracker(client, item) for item in trackables]
 

homeassistant/components/tractive/diagnostics.py

@@ -5,20 +5,19 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_EMAIL, CONF_PASSWORD
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN, TRACKABLES
+from . import TractiveConfigEntry
 
 TO_REDACT = {CONF_PASSWORD, CONF_EMAIL, "title", "_id"}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: TractiveConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    trackables = hass.data[DOMAIN][config_entry.entry_id][TRACKABLES]
+    trackables = config_entry.runtime_data.trackables
 
     return async_redact_data(
         {

homeassistant/components/tractive/sensor.py

@@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_BATTERY_LEVEL,
     PERCENTAGE,
@@ -23,7 +22,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from . import Trackables, TractiveClient
+from . import Trackables, TractiveClient, TractiveConfigEntry
 from .const import (
     ATTR_ACTIVITY_LABEL,
     ATTR_CALORIES,
@@ -34,9 +33,6 @@ from .const import (
     ATTR_MINUTES_REST,
     ATTR_SLEEP_LABEL,
     ATTR_TRACKER_STATE,
-    CLIENT,
-    DOMAIN,
-    TRACKABLES,
     TRACKER_HARDWARE_STATUS_UPDATED,
     TRACKER_WELLNESS_STATUS_UPDATED,
 )
@@ -183,11 +179,13 @@ SENSOR_TYPES: tuple[TractiveSensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: TractiveConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Tractive device trackers."""
-    client = hass.data[DOMAIN][entry.entry_id][CLIENT]
-    trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
+    client = entry.runtime_data.client
+    trackables = entry.runtime_data.trackables
 
     entities = [
         TractiveSensor(client, item, description)

homeassistant/components/tractive/switch.py

@@ -9,19 +9,15 @@ from typing import Any, Literal, cast
 from aiotractive.exceptions import TractiveError
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import Trackables, TractiveClient
+from . import Trackables, TractiveClient, TractiveConfigEntry
 from .const import (
     ATTR_BUZZER,
     ATTR_LED,
     ATTR_LIVE_TRACKING,
-    CLIENT,
-    DOMAIN,
-    TRACKABLES,
     TRACKER_SWITCH_STATUS_UPDATED,
 )
 from .entity import TractiveEntity
@@ -59,11 +55,13 @@ SWITCH_TYPES: tuple[TractiveSwitchEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: TractiveConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Tractive switches."""
-    client = hass.data[DOMAIN][entry.entry_id][CLIENT]
-    trackables = hass.data[DOMAIN][entry.entry_id][TRACKABLES]
+    client = entry.runtime_data.client
+    trackables = entry.runtime_data.trackables
 
     entities = [
         TractiveSwitch(client, item, description)

tests/components/tractive/conftest.py

@@ -0,0 +1,53 @@
+"""Common fixtures for the Tractive tests."""
+
+from collections.abc import Generator
+from unittest.mock import AsyncMock, Mock, patch
+
+from aiotractive.trackable_object import TrackableObject
+from aiotractive.tracker import Tracker
+import pytest
+
+from homeassistant.components.tractive.const import DOMAIN
+from homeassistant.const import CONF_EMAIL, CONF_PASSWORD
+
+from tests.common import MockConfigEntry, load_json_object_fixture
+
+
+@pytest.fixture
+def mock_tractive_client() -> Generator[AsyncMock, None, None]:
+    """Mock a Tractive client."""
+
+    trackable_object = load_json_object_fixture("tractive/trackable_object.json")
+    with (
+        patch(
+            "homeassistant.components.tractive.aiotractive.Tractive", autospec=True
+        ) as mock_client,
+    ):
+        client = mock_client.return_value
+        client.authenticate.return_value = {"user_id": "12345"}
+        client.trackable_objects.return_value = [
+            Mock(
+                spec=TrackableObject,
+                _id="xyz123",
+                type="pet",
+                details=AsyncMock(return_value=trackable_object),
+            ),
+        ]
+        client.tracker.return_value = Mock(spec=Tracker)
+
+        yield client
+
+
+@pytest.fixture
+def mock_config_entry() -> MockConfigEntry:
+    """Mock a config entry."""
+    return MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_EMAIL: "test-email@example.com",
+            CONF_PASSWORD: "test-password",
+        },
+        unique_id="very_unique_string",
+        entry_id="3bd2acb0e4f0476d40865546d0d91921",
+        title="Test Pet",
+    )

tests/components/tractive/fixtures/trackable_object.json

@@ -0,0 +1,42 @@
+{
+  "device_id": "54321",
+  "details": {
+    "_id": "xyz123",
+    "_version": "123abc",
+    "name": "Test Pet",
+    "pet_type": "DOG",
+    "breed_ids": [],
+    "gender": "F",
+    "birthday": 1572606592,
+    "profile_picture_frame": null,
+    "height": 0.56,
+    "length": null,
+    "weight": 23700,
+    "chip_id": "",
+    "neutered": true,
+    "personality": [],
+    "lost_or_dead": null,
+    "lim": null,
+    "ribcage": null,
+    "weight_is_default": null,
+    "height_is_default": null,
+    "birthday_is_default": null,
+    "breed_is_default": null,
+    "instagram_username": "",
+    "profile_picture_id": null,
+    "cover_picture_id": null,
+    "characteristic_ids": [],
+    "gallery_picture_ids": [],
+    "activity_settings": {
+      "_id": "345abc",
+      "_version": "ccaabb4",
+      "daily_goal": 1000,
+      "daily_distance_goal": 2000,
+      "daily_active_minutes_goal": 120,
+      "activity_category_thresholds_override": null,
+      "_type": "activity_setting"
+    },
+    "_type": "pet_detail",
+    "read_only": false
+  }
+}

tests/components/tractive/snapshots/test_diagnostics.ambr

@@ -0,0 +1,71 @@
+# serializer version: 1
+# name: test_entry_diagnostics
+  dict({
+    'config_entry': dict({
+      'data': dict({
+        'email': '**REDACTED**',
+        'password': '**REDACTED**',
+      }),
+      'disabled_by': None,
+      'domain': 'tractive',
+      'entry_id': '3bd2acb0e4f0476d40865546d0d91921',
+      'minor_version': 1,
+      'options': dict({
+      }),
+      'pref_disable_new_entities': False,
+      'pref_disable_polling': False,
+      'source': 'user',
+      'title': '**REDACTED**',
+      'unique_id': 'very_unique_string',
+      'version': 1,
+    }),
+    'trackables': list([
+      dict({
+        'details': dict({
+          '_id': '**REDACTED**',
+          '_type': 'pet_detail',
+          '_version': '123abc',
+          'activity_settings': dict({
+            '_id': '**REDACTED**',
+            '_type': 'activity_setting',
+            '_version': 'ccaabb4',
+            'activity_category_thresholds_override': None,
+            'daily_active_minutes_goal': 120,
+            'daily_distance_goal': 2000,
+            'daily_goal': 1000,
+          }),
+          'birthday': 1572606592,
+          'birthday_is_default': None,
+          'breed_ids': list([
+          ]),
+          'breed_is_default': None,
+          'characteristic_ids': list([
+          ]),
+          'chip_id': '',
+          'cover_picture_id': None,
+          'gallery_picture_ids': list([
+          ]),
+          'gender': 'F',
+          'height': 0.56,
+          'height_is_default': None,
+          'instagram_username': '',
+          'length': None,
+          'lim': None,
+          'lost_or_dead': None,
+          'name': 'Test Pet',
+          'neutered': True,
+          'personality': list([
+          ]),
+          'pet_type': 'DOG',
+          'profile_picture_frame': None,
+          'profile_picture_id': None,
+          'read_only': False,
+          'ribcage': None,
+          'weight': 23700,
+          'weight_is_default': None,
+        }),
+        'device_id': '54321',
+      }),
+    ]),
+  })
+# ---

tests/components/tractive/test_diagnostics.py

@@ -0,0 +1,31 @@
+"""Test the Tractive diagnostics."""
+
+from unittest.mock import AsyncMock, patch
+
+from syrupy import SnapshotAssertion
+
+from homeassistant.components.tractive.const import DOMAIN
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from tests.common import MockConfigEntry
+from tests.components.diagnostics import get_diagnostics_for_config_entry
+from tests.typing import ClientSessionGenerator
+
+
+async def test_entry_diagnostics(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    snapshot: SnapshotAssertion,
+    mock_tractive_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test config entry diagnostics."""
+    mock_config_entry.add_to_hass(hass)
+    with patch("homeassistant.components.tractive.PLATFORMS", []):
+        assert await async_setup_component(hass, DOMAIN, {})
+    result = await get_diagnostics_for_config_entry(
+        hass, hass_client, mock_config_entry
+    )
+
+    assert result == snapshot