From e62c9d338e8ca73b94665ed603fd1b5b13ea5fc4 Mon Sep 17 00:00:00 2001 From: Maciej Bieniek Date: Fri, 24 Sep 2021 08:45:03 +0200 Subject: [PATCH] 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 | 51 +++++++++++++++++-- homeassistant/components/tractive/const.py | 3 ++ .../components/tractive/device_tracker.py | 35 ++++++------- homeassistant/components/tractive/sensor.py | 23 +++------ 4 files changed, 74 insertions(+), 38 deletions(-) diff --git a/homeassistant/components/tractive/__init__.py b/homeassistant/components/tractive/__init__.py index 60014852895..c380471769c 100644 --- a/homeassistant/components/tractive/__init__.py +++ b/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 diff --git a/homeassistant/components/tractive/const.py b/homeassistant/components/tractive/const.py index cb525d538e4..7f1b5ddb4f2 100644 --- a/homeassistant/components/tractive/const.py +++ b/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" diff --git a/homeassistant/components/tractive/device_tracker.py b/homeassistant/components/tractive/device_tracker.py index c1652c27b8f..1e35e41fc8a 100644 --- a/homeassistant/components/tractive/device_tracker.py +++ b/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.""" diff --git a/homeassistant/components/tractive/sensor.py b/homeassistant/components/tractive/sensor.py index ba2f330f894..9fd8ee6ac5f 100644 --- a/homeassistant/components/tractive/sensor.py +++ b/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)