Merge branch 'development'

Author: Crewski

custom_components/pumpspy_ha/__init__.py

@@ -2,15 +2,13 @@
 from __future__ import annotations
 from datetime import timedelta
 import logging
-import async_timeout
 
 from homeassistant.helpers.device_registry import DeviceEntry
 
-from .pumpspy import Pumpspy
+from .pypumpspy import InvalidAccessToken, Pumpspy
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
-    CONF_ACCESS_TOKEN,
     CONF_PASSWORD,
     CONF_USERNAME,
     Platform,
@@ -20,35 +18,26 @@
     DataUpdateCoordinator,
 )
 
-# from homeassistant.exceptions import ConfigEntryAuthFailed
-
 
 from .const import (
-    CONF_DEVICE_NAME,
     CONF_DEVICEID,
-    CONF_REFRESH_TOKEN,
     DOMAIN,
-    MANUFACTURER,
 )
 
-# TODO List the platforms that you want to support.
-# For your initial PR, limit it to 1 platform.
+
 PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.BINARY_SENSOR]
 
 _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Pumpspy-HA from a config entry."""
-    api = Pumpspy(hass)
-    api.set_variables(
-        entry.data[CONF_ACCESS_TOKEN],
-        entry.data[CONF_REFRESH_TOKEN],
-        entry.data[CONF_DEVICEID],
-        entry.data[CONF_USERNAME],
-        entry.data[CONF_DEVICE_NAME],
-        entry.data[CONF_PASSWORD],
+    api: Pumpspy = Pumpspy(
+        username=entry.data[CONF_USERNAME],
+        password=entry.data[CONF_PASSWORD],
+        device_id=entry.data[CONF_DEVICEID],
     )
+    await api.setup()
     coordinator = PumpspyCoordinator(hass, api)
     await coordinator.async_config_entry_first_refresh()
 
@@ -89,31 +78,10 @@ def __init__(self, hass, api):
         self.api = api
 
     async def _async_update_data(self):
-        """Fetch data from API endpoint.
-
-        This is the place to pre-process the data to lookup tables
-        so entities can quickly look up their data.
-        """
+        """Fetch data from API endpoint."""
         try:
-            # Note: asyncio.TimeoutError and aiohttp.ClientError are already
-            # handled by the data update coordinator.
-            data = {}
-            async with async_timeout.timeout(30):
-                data["current"] = await self.api.fetch_current_data()
-            async with async_timeout.timeout(30):
-                data["main_daily"] = await self.api.fetch_interval_data("ac", "day")
-            async with async_timeout.timeout(30):
-                data["backup_daily"] = await self.api.fetch_interval_data("dc", "day")
-            async with async_timeout.timeout(30):
-                data["main_monthly"] = await self.api.fetch_interval_data("ac", "month")
-            async with async_timeout.timeout(30):
-                data["backup_monthly"] = await self.api.fetch_interval_data(
-                    "dc", "month"
-                )
-            async with async_timeout.timeout(30):
-                data["main_weekly"] = await self.api.fetch_interval_data("ac", "week")
-            async with async_timeout.timeout(30):
-                data["backup_weekly"] = await self.api.fetch_interval_data("dc", "week")
-            return data
-        except Exception as err:
-            self.logger.error(err)
+            return await self.api.fetch_data()
+        except InvalidAccessToken:
+            _LOGGER.info("Access token expired, will try again")
+        except ConnectionError as err:
+            _LOGGER.error(err)

custom_components/pumpspy_ha/binary_sensor.py

@@ -6,8 +6,6 @@
 from homeassistant.helpers.entity import EntityCategory
 
 from .entity import PumpspyEntity
-from homeassistant.core import callback
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
@@ -34,16 +32,18 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     coordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     new_devices = [
-        AlertBinarySensor(coordinator, ALERT_CONNECTED),
-        AlertBinarySensor(coordinator, ALERT_HIGH_WATER),
-        AlertBinarySensor(coordinator, ALERT_AC_POWER_LOSS),
-        AlertBinarySensor(coordinator, ALERT_EXCESSIVE_CURRENT),
-        AlertBinarySensor(coordinator, ALERT_EXCESSIVE_RUN_TIME),
-        AlertBinarySensor(coordinator, ALERT_BATTERY_CHARGE_LEVEL),
-        AlertBinarySensor(coordinator, ALERT_BACKUP_EXCESSIVE_CURRET),
-        AlertBinarySensor(coordinator, ALERT_BACKUP_EXCESSIVE_RUN_TIME),
-        AlertBinarySensor(coordinator, ALERT_PRIMARY_PUMP_FAILURE),
-        AlertBinarySensor(coordinator, ALERT_BACKUP_PUMP_FAILURE),
+        AlertBinarySensor(coordinator=coordinator, alert=ALERT_CONNECTED),
+        AlertBinarySensor(coordinator=coordinator, alert=ALERT_HIGH_WATER),
+        AlertBinarySensor(coordinator=coordinator, alert=ALERT_AC_POWER_LOSS),
+        AlertBinarySensor(coordinator=coordinator, alert=ALERT_EXCESSIVE_CURRENT),
+        AlertBinarySensor(coordinator=coordinator, alert=ALERT_EXCESSIVE_RUN_TIME),
+        AlertBinarySensor(coordinator=coordinator, alert=ALERT_BATTERY_CHARGE_LEVEL),
+        AlertBinarySensor(coordinator=coordinator, alert=ALERT_BACKUP_EXCESSIVE_CURRET),
+        AlertBinarySensor(
+            coordinator=coordinator, alert=ALERT_BACKUP_EXCESSIVE_RUN_TIME
+        ),
+        AlertBinarySensor(coordinator=coordinator, alert=ALERT_PRIMARY_PUMP_FAILURE),
+        AlertBinarySensor(coordinator=coordinator, alert=ALERT_BACKUP_PUMP_FAILURE),
     ]
     if new_devices:
         async_add_entities(new_devices)
@@ -54,9 +54,7 @@ class AlertBinarySensor(PumpspyEntity, BinarySensorEntity):
 
     def __init__(self, coordinator, alert: str):
         """Initialize the sensor."""
-        # super().__init__(coordinator)
-        self.coordinator = coordinator
-        self.coordinator_context = object()
+        super().__init__(coordinator=coordinator)
         self._available = True
         self._alert = alert
         self._attr_device_class = (
@@ -73,20 +71,6 @@ def __init__(self, coordinator, alert: str):
             f'{device_info[CONF_DEVICE_NAME]} {alert.replace("_", " ").title()}'
         )
 
-    # @callback
-    # def _handle_coordinator_update(self) -> None:
-    #     """Handle updated data from the coordinator."""
-    #     if self.coordinator.data is None:
-    #         return
-    #     val = self.coordinator.data["current"][0][self._alert]["state"]
-    #     if self._alert == ALERT_BATTERY_CHARGE_LEVEL:
-    #         val = not bool(val)
-    #     self._attr_is_on = val
-    #     self._attr_extra_state_attributes = {
-    #         "message": self.coordinator.data["current"][0][self._alert]["message"]
-    #     }
-    #     self.async_write_ha_state()
-
     @property
     def is_on(self) -> bool | None:
         val = self.coordinator.data["current"][0][self._alert]["state"]

custom_components/pumpspy_ha/config_flow.py

@@ -5,21 +5,19 @@
 from typing import Any
 
 import voluptuous as vol
+from .pypumpspy import Pumpspy
 
 from homeassistant import config_entries
-from .pumpspy import Pumpspy
 from homeassistant.data_entry_flow import FlowResult
-from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import selector
 
 from homeassistant.const import (
-    CONF_DEVICE_ID,
     CONF_USERNAME,
-    CONF_ACCESS_TOKEN,
     CONF_PASSWORD,
 )
 
 from .const import (
+    CONF_DEVICEID,
     DOMAIN,
 )
 
@@ -29,23 +27,28 @@
 class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Pumpspy-HA."""
 
-    VERSION = 1
+    VERSION = 2
 
-    pumpspy: Pumpspy
+    pumpspy: Pumpspy = None
     locations = None
     devices = None
 
+    data: dict[str, Any] = {}
+
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
         """Handle the initial step."""
 
         errors = {}
         if user_input is not None:
-            self.pumpspy = Pumpspy(self.hass)
-            self.locations = await self.pumpspy.get_locations(
-                user_input[CONF_USERNAME], user_input[CONF_PASSWORD]
+            self.pumpspy = Pumpspy(
+                username=user_input[CONF_USERNAME], password=user_input[CONF_PASSWORD]
             )
+            await self.pumpspy.setup()
+            self.data[CONF_USERNAME] = user_input[CONF_USERNAME]
+            self.data[CONF_PASSWORD] = user_input[CONF_PASSWORD]
+            self.locations = await self.pumpspy.get_locations()
             if self.locations:
                 return await self.async_step_location()
 
@@ -67,12 +70,14 @@ async def async_step_location(
 
         # skip this step if there is only 1 location
         if len(self.locations) == 1:
-            self.devices = await self.pumpspy.get_devices(self.locations[0]["lid"])
+            self.pumpspy.set_location(self.locations[0]["lid"])
+            self.devices = await self.pumpspy.get_devices()
             if self.devices:
                 return await self.async_step_device()
 
         if user_input is not None:
-            self.devices = await self.pumpspy.get_devices(user_input["location"])
+            self.pumpspy.set_location(self.locations[0]["lid"])
+            self.devices = await self.pumpspy.get_devices()
             if self.devices:
                 return await self.async_step_device()
 
@@ -103,14 +108,12 @@ async def async_step_device(
 
         # skip this step if there is only 1 device
         if len(self.devices) == 1:
-            self.pumpspy.set_device_info(
-                self.devices[0]["deviceid"], self.devices[0]["device_types_name"]
-            )
+            self.data[CONF_DEVICEID] = self.devices[0]["deviceid"]
             await self.async_set_unique_id(self.devices[0]["deviceid"])
             self._abort_if_unique_id_configured()
             return self.async_create_entry(
                 title=f'Pumpspy ({self.devices[0]["device_types_name"]})',
-                data=self.pumpspy.get_variables(),
+                data=self.data,
             )
 
         if user_input is not None:
@@ -119,12 +122,11 @@ async def async_step_device(
             device_name = "Unknown"
 
             for x in self.devices:
-                if x['deviceid'] == int(user_input['device']):
-                    device_name = x['device_types_name']
-                    
-            self.pumpspy.set_device_info(user_input["device"], device_name)
+                if x["deviceid"] == int(user_input["device"]):
+                    device_name = x["device_types_name"]
+            self.data[CONF_DEVICEID] = user_input["device"]
             return self.async_create_entry(
-                title=f"Pumpspy ({device_name})", data=self.pumpspy.get_variables()
+                title=f"Pumpspy ({device_name})", data=self.data
             )
 
         options = []

custom_components/pumpspy_ha/entity.py

@@ -10,8 +10,9 @@
 class PumpspyEntity(CoordinatorEntity[PumpspyCoordinator]):
     """Defines a base Pumpspy entity."""
 
-    def __init__(self, *, coordinator: PumpspyCoordinator) -> None:
+    def __init__(self, coordinator: PumpspyCoordinator) -> None:
         """Initialize the entity."""
+        self.coordinator = coordinator
         super().__init__(coordinator)
 
     @property