Extracted data update coordinator.

custom_components/mygekko/__init__.py

@@ -6,37 +6,19 @@
 """
 import asyncio
 import logging
-from datetime import timedelta
 
+from custom_components.mygekko.coordinator import MyGekkoDataUpdateCoordinator
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY
-from homeassistant.const import CONF_IP_ADDRESS
-from homeassistant.const import CONF_PASSWORD
-from homeassistant.const import CONF_USERNAME
 from homeassistant.core import Config
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryError
-from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
-from homeassistant.helpers.update_coordinator import UpdateFailed
-from PyMyGekko import MyGekkoApiClientBase
-from PyMyGekko import MyGekkoDemoModeClient
-from PyMyGekko import MyGekkoLocalApiClient
-from PyMyGekko import MyGekkoQueryApiClient
-
-from .const import CONF_CONNECTION_DEMO_MODE
-from .const import CONF_CONNECTION_LOCAL
+
 from .const import CONF_CONNECTION_MY_GEKKO_CLOUD
 from .const import CONF_CONNECTION_TYPE
-from .const import CONF_GEKKOID
 from .const import DOMAIN
 from .const import PLATFORMS
 from .const import STARTUP_MESSAGE
 
-
-SCAN_INTERVAL = timedelta(seconds=30)
-
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
@@ -53,48 +35,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
 
     _LOGGER.debug("async_setup_entry %s", entry.data.get(CONF_CONNECTION_TYPE))
 
-    client = None
-
-    if entry.data.get(CONF_CONNECTION_TYPE) == CONF_CONNECTION_MY_GEKKO_CLOUD:
-        username = entry.data.get(CONF_USERNAME)
-        apikey = entry.data.get(CONF_API_KEY)
-        gekkoid = entry.data.get(CONF_GEKKOID)
-
-        session = async_get_clientsession(hass)
-
-        client = MyGekkoQueryApiClient(username, apikey, gekkoid, session)
-
-    if entry.data.get(CONF_CONNECTION_TYPE) == CONF_CONNECTION_LOCAL:
-        username = entry.data.get(CONF_USERNAME)
-        password = entry.data.get(CONF_PASSWORD)
-        ip_address = entry.data.get(CONF_IP_ADDRESS)
-
-        session = async_get_clientsession(hass, verify_ssl=False)
+    coordinator = MyGekkoDataUpdateCoordinator(hass, entry=entry)
 
-        client = MyGekkoLocalApiClient(username, password, session, ip_address)
-
-    if entry.data.get(CONF_CONNECTION_TYPE) == CONF_CONNECTION_DEMO_MODE:
-        client = MyGekkoDemoModeClient()
-
-    if client is None:
-        _LOGGER.exception("async_refresh failed: client is None")
-        raise ConfigEntryError
-
-    coordinator = MyGekkoDataUpdateCoordinator(hass, client=client)
-    await coordinator.async_refresh()
-
-    if not coordinator.last_update_success:
-        _LOGGER.exception("async_refresh failed")
-        raise ConfigEntryNotReady
+    await coordinator.async_config_entry_first_refresh()
 
     hass.data[DOMAIN][entry.entry_id] = coordinator
 
     for platform in PLATFORMS:
-        if entry.options.get(platform, True):
-            coordinator.platforms.append(platform)
-            hass.async_add_job(
-                hass.config_entries.async_forward_entry_setup(entry, platform)
-            )
+        coordinator.platforms.append(platform)
+        hass.async_add_job(
+            hass.config_entries.async_forward_entry_setup(entry, platform)
+        )
 
     entry.add_update_listener(async_reload_entry)
     return True
@@ -124,30 +75,6 @@ async def async_migrate_entry(hass, config_entry: ConfigEntry):
     return True
 
 
-class MyGekkoDataUpdateCoordinator(DataUpdateCoordinator):
-    """Class to manage fetching data from the API."""
-
-    def __init__(
-        self,
-        hass: HomeAssistant,
-        client: MyGekkoApiClientBase,
-    ) -> None:
-        """Initialize."""
-        self.api = client
-        self.platforms = []
-
-        super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL)
-
-    async def _async_update_data(self):
-        """Update data via library."""
-        _LOGGER.debug("_async_update_data ")
-        try:
-            return await self.api.read_data()
-        except Exception as exception:
-            _LOGGER.exception("_async_update_data failed")
-            raise UpdateFailed() from exception
-
-
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Handle removal of an entry."""
     coordinator = hass.data[DOMAIN][entry.entry_id]

custom_components/mygekko/coordinator.py

@@ -0,0 +1,80 @@
+import logging
+from datetime import timedelta
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_API_KEY
+from homeassistant.const import CONF_IP_ADDRESS
+from homeassistant.const import CONF_PASSWORD
+from homeassistant.const import CONF_USERNAME
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryError
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+from homeassistant.helpers.update_coordinator import UpdateFailed
+from PyMyGekko import MyGekkoDemoModeClient
+from PyMyGekko import MyGekkoLocalApiClient
+from PyMyGekko import MyGekkoQueryApiClient
+
+from .const import CONF_CONNECTION_DEMO_MODE
+from .const import CONF_CONNECTION_LOCAL
+from .const import CONF_CONNECTION_MY_GEKKO_CLOUD
+from .const import CONF_CONNECTION_TYPE
+from .const import CONF_GEKKOID
+from .const import DOMAIN
+
+
+SCAN_INTERVAL = timedelta(seconds=30)
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+class MyGekkoDataUpdateCoordinator(DataUpdateCoordinator):
+    """Class to manage fetching data from the API."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        entry: ConfigEntry,
+    ) -> None:
+        """Initialize."""
+        self.platforms = []
+
+        super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL)
+
+        client = None
+
+        if entry.data.get(CONF_CONNECTION_TYPE) == CONF_CONNECTION_MY_GEKKO_CLOUD:
+            username = entry.data.get(CONF_USERNAME)
+            apikey = entry.data.get(CONF_API_KEY)
+            gekkoid = entry.data.get(CONF_GEKKOID)
+
+            session = async_get_clientsession(hass)
+            client = MyGekkoQueryApiClient(username, apikey, gekkoid, session)
+
+        if entry.data.get(CONF_CONNECTION_TYPE) == CONF_CONNECTION_LOCAL:
+            username = entry.data.get(CONF_USERNAME)
+            password = entry.data.get(CONF_PASSWORD)
+            ip_address = entry.data.get(CONF_IP_ADDRESS)
+
+            session = async_get_clientsession(hass, verify_ssl=False)
+            client = MyGekkoLocalApiClient(username, password, session, ip_address)
+
+        if entry.data.get(CONF_CONNECTION_TYPE) == CONF_CONNECTION_DEMO_MODE:
+            client = MyGekkoDemoModeClient()
+
+        if client is None:
+            _LOGGER.exception(
+                "Creating MyGekkoDataUpdateCoordinator failed: client is None"
+            )
+            raise ConfigEntryError
+
+        self.api = client
+
+    async def _async_update_data(self):
+        """Update data via library."""
+        _LOGGER.debug("_async_update_data ")
+        try:
+            return await self.api.read_data()
+        except Exception as exception:
+            _LOGGER.exception("_async_update_data failed")
+            raise UpdateFailed() from exception