use HA DataUpdateCoordinator

custom_components/trakt/__init__.py

@@ -23,22 +23,21 @@
     CONF_SCAN_INTERVAL,
 )
 from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow
-from homeassistant.helpers.dispatcher import async_dispatcher_send
-from homeassistant.helpers.event import async_track_time_interval
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 from trakt.calendar import MyShowCalendar
 from trakt.tv import TVShow
 
 from . import config_flow
 from .const import (
     CARD_DEFAULT,
     CONF_DAYS,
-    DATA_UPDATED,
     DEFAULT_DAYS,
     DEFAULT_SCAN_INTERVAL,
     DOMAIN,
     OAUTH2_AUTHORIZE,
     OAUTH2_TOKEN,
 )
+from homeassistant.exceptions import ConfigEntryNotReady
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -71,15 +70,16 @@ async def async_setup_entry(hass, entry) -> bool:
     session = config_entry_oauth2_flow.OAuth2Session(hass, entry, implementation)
     await session.async_ensure_token_valid()
 
-    trakt_data = Trakt_Data(hass, entry, session)
-
-    if not await trakt_data.async_setup():
+    coordinator = Trakt_Data(hass, entry, session)
+    if not await coordinator.async_setup():
         return False
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = trakt_data
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+
     hass.async_create_task(
         hass.config_entries.async_forward_entry_setup(entry, "sensor")
     )
+
     return True
 
 
@@ -95,7 +95,7 @@ async def async_unload_entry(hass, entry) -> bool:
     return True
 
 
-class Trakt_Data:
+class Trakt_Data(DataUpdateCoordinator):
     """Represent Trakt data."""
 
     def __init__(
@@ -109,8 +109,19 @@ def __init__(
         self.config_entry = config_entry
         self.session = session
         self.unsub_timer = None
-        self.details = {}
-        self.calendar = None
+        self.calendar = {}
+
+        super().__init__(
+            self.hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_method=self.async_update,
+            update_interval=timedelta(
+                minutes=self.config_entry.options.get(
+                    CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL
+                )
+            ),
+        )
 
     @property
     def days(self):
@@ -135,15 +146,12 @@ async def async_update(self, *_):
             self.calendar = await self.hass.async_add_executor_job(
                 MyShowCalendar, {CONF_DAYS: self.days}
             )
-        except trakt.errors.OAuthException:
-            _LOGGER.error(
-                "Trakt api is unauthrized. Please remove the entry and reconfigure."
-            )
-            return
+        except trakt.errors.TraktInternalException:
+            _LOGGER.error("Trakt api encountered an internal error.")
+            raise UpdateFailed
 
         if not self.calendar:
             _LOGGER.warning("Trakt upcoming calendar is empty")
-            return
 
         for show in self.calendar:
             if not show or show.show in self.exclude:
@@ -203,45 +211,34 @@ async def async_update(self, *_):
             }
             card_json.append(card_item)
 
-        self.details = json.dumps(card_json)
-
-        _LOGGER.debug("Trakt data updated")
-        async_dispatcher_send(self.hass, DATA_UPDATED)
+        return json.dumps(card_json)
 
     async def async_setup(self):
         """Set up Trakt Data."""
         trakt.core.OAUTH_TOKEN = self.session.token[CONF_ACCESS_TOKEN]
         trakt.core.CLIENT_ID = self.config_entry.data[CONF_CLIENT_ID]
 
         try:
-            await self.async_update()
+            await self.async_refresh()
         except trakt.errors.OAuthException:
             _LOGGER.error(
                 "Trakt api is unauthrized. Please remove the entry and reconfigure."
             )
             return False
 
-        await self.async_set_scan_interval(self.scan_interval)
-        self.config_entry.add_update_listener(self.async_options_updated)
+        if not self.last_update_success:
+            raise ConfigEntryNotReady
+        self.config_entry.add_update_listener(async_options_updated)
 
         return True
 
-    async def async_set_scan_interval(self, scan_interval):
-        """Update scan interval."""
 
-        if self.unsub_timer is not None:
-            self.unsub_timer()
-        self.unsub_timer = async_track_time_interval(
-            self.hass, self.async_update, timedelta(minutes=scan_interval)
-        )
-        await self.async_update()
-
-    @staticmethod
-    async def async_options_updated(hass, entry):
-        """Triggered by config entry options updates."""
-        await hass.data[DOMAIN][entry.entry_id].async_set_scan_interval(
-            entry.options[CONF_SCAN_INTERVAL]
-        )
+async def async_options_updated(hass, entry):
+    """Triggered by config entry options updates."""
+    hass.data[DOMAIN][entry.entry_id].update_interval = timedelta(
+        minutes=entry.options[CONF_SCAN_INTERVAL]
+    )
+    await hass.data[DOMAIN][entry.entry_id].async_request_refresh()
 
 
 def days_until(date):

custom_components/trakt/sensor.py

@@ -1,25 +1,24 @@
 """Sensor platform for Trakt"""
 from homeassistant.const import ATTR_ATTRIBUTION, CONF_NAME, CONF_CLIENT_ID
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
+
 from homeassistant.helpers.entity import Entity
 
-from .const import ATTRIBUTION, DATA_UPDATED, DOMAIN
+from .const import ATTRIBUTION, DOMAIN
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up device tracker for Mikrotik component."""
-    tk_data = hass.data[DOMAIN][config_entry.entry_id]
-
-    async_add_entities([TraktUpcomingCalendarSensor(tk_data)], True)
+    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    async_add_entities([TraktUpcomingCalendarSensor(coordinator)], True)
 
 
 class TraktUpcomingCalendarSensor(Entity):
     """Representation of a Trakt Upcoming Calendar sensor."""
 
-    def __init__(self, tk_data):
+    def __init__(self, coordinator):
         """Initialize the sensor."""
-        self.tk_data = tk_data
-        self._name = tk_data.config_entry.data[CONF_NAME]
+        self.coordinator = coordinator
+        self._name = coordinator.config_entry.data[CONF_NAME]
 
     @property
     def name(self):
@@ -29,7 +28,7 @@ def name(self):
     @property
     def unique_id(self):
         """Return the unique id of the entity."""
-        return self.tk_data.config_entry.data[CONF_CLIENT_ID]
+        return self.coordinator.config_entry.data[CONF_CLIENT_ID]
 
     @property
     def should_poll(self):
@@ -39,7 +38,7 @@ def should_poll(self):
     @property
     def state(self):
         """Return the state of the sensor."""
-        return len(self.tk_data.calendar)
+        return len(self.coordinator.calendar) if self.coordinator.calendar else 0
 
     @property
     def icon(self):
@@ -54,12 +53,19 @@ def unit_of_measurement(self):
     @property
     def device_state_attributes(self):
         """Return the state attributes of the sensor."""
-        attributes = {"data": self.tk_data.details, ATTR_ATTRIBUTION: ATTRIBUTION}
+        attributes = {
+            "data": self.coordinator.data,
+            ATTR_ATTRIBUTION: ATTRIBUTION,
+        }
 
         return attributes
 
     async def async_added_to_hass(self):
         """Handle entity which will be added."""
         self.async_on_remove(
-            async_dispatcher_connect(self.hass, DATA_UPDATED, self.async_write_ha_state)
+            self.coordinator.async_add_listener(self.async_write_ha_state)
         )
+
+    async def async_update(self):
+        """Request coordinator to update data."""
+        await self.coordinator.async_request_refresh()