Merge pull request #3 from Crewski/0.1.3

Add access token refresh and HA device

Author: Crewski

custom_components/pumpspy_ha/__init__.py

@@ -4,18 +4,32 @@
 import logging
 import async_timeout
 
+from homeassistant.helpers.device_registry import DeviceEntry
+
 from .pumpspy import Pumpspy
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_ACCESS_TOKEN, CONF_USERNAME, Platform
+from homeassistant.const import (
+    CONF_ACCESS_TOKEN,
+    CONF_PASSWORD,
+    CONF_USERNAME,
+    Platform,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import (
     DataUpdateCoordinator,
 )
 
 # from homeassistant.exceptions import ConfigEntryAuthFailed
 
-from .const import CONF_DEVICE_NAME, CONF_DEVICEID, CONF_REFRESH_TOKEN, DOMAIN
+
+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.
@@ -33,13 +47,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         entry.data[CONF_DEVICEID],
         entry.data[CONF_USERNAME],
         entry.data[CONF_DEVICE_NAME],
+        entry.data[CONF_PASSWORD],
     )
     coordinator = PumpspyCoordinator(hass, api)
+    await coordinator.async_config_entry_first_refresh()
+
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
-    await coordinator.async_config_entry_first_refresh()
-
     return True
 
 
@@ -51,6 +66,13 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return unload_ok
 
 
+async def async_remove_config_entry_device(
+    hass: HomeAssistant, config_entry: ConfigEntry, device_entry: DeviceEntry
+) -> bool:
+    """Remove a config entry from a device."""
+    return True
+
+
 class PumpspyCoordinator(DataUpdateCoordinator):
     """Pumpspy coordinator."""
 
@@ -76,21 +98,21 @@ async def _async_update_data(self):
             # Note: asyncio.TimeoutError and aiohttp.ClientError are already
             # handled by the data update coordinator.
             data = {}
-            async with async_timeout.timeout(10):
+            async with async_timeout.timeout(30):
                 data["current"] = await self.api.fetch_current_data()
-            async with async_timeout.timeout(10):
+            async with async_timeout.timeout(30):
                 data["main_daily"] = await self.api.fetch_interval_data("ac", "day")
-            async with async_timeout.timeout(10):
+            async with async_timeout.timeout(30):
                 data["backup_daily"] = await self.api.fetch_interval_data("dc", "day")
-            async with async_timeout.timeout(10):
+            async with async_timeout.timeout(30):
                 data["main_monthly"] = await self.api.fetch_interval_data("ac", "month")
-            async with async_timeout.timeout(10):
+            async with async_timeout.timeout(30):
                 data["backup_monthly"] = await self.api.fetch_interval_data(
                     "dc", "month"
                 )
-            async with async_timeout.timeout(10):
+            async with async_timeout.timeout(30):
                 data["main_weekly"] = await self.api.fetch_interval_data("ac", "week")
-            async with async_timeout.timeout(10):
+            async with async_timeout.timeout(30):
                 data["backup_weekly"] = await self.api.fetch_interval_data("dc", "week")
             return data
         except Exception as err:

custom_components/pumpspy_ha/binary_sensor.py

@@ -1,6 +1,11 @@
 """Platform for binary sensor integration."""
+from __future__ import annotations
+from typing import Any
+from collections.abc import Mapping
 
+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 (
@@ -44,19 +49,22 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         async_add_entities(new_devices)
 
 
-class AlertBinarySensor(CoordinatorEntity, BinarySensorEntity):
+class AlertBinarySensor(PumpspyEntity, BinarySensorEntity):
     """Alert Binary Sensor"""
 
     def __init__(self, coordinator, alert: str):
         """Initialize the sensor."""
-        super().__init__(coordinator)
+        # super().__init__(coordinator)
+        self.coordinator = coordinator
+        self.coordinator_context = object()
         self._available = True
         self._alert = alert
         self._attr_device_class = (
             BinarySensorDeviceClass.CONNECTIVITY
             if alert == ALERT_CONNECTED
             else BinarySensorDeviceClass.PROBLEM
         )
+        self._attr_entity_category = EntityCategory.DIAGNOSTIC
 
         device_info = self.coordinator.api.get_device_info()
 
@@ -65,16 +73,27 @@ 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
+    # @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"]
         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()
+        return val
+
+    @property
+    def extra_state_attributes(self) -> Mapping[str, Any] | None:
+        return {"message": self.coordinator.data["current"][0][self._alert]["message"]}

custom_components/pumpspy_ha/config_flow.py

@@ -13,7 +13,9 @@
 from homeassistant.helpers import selector
 
 from homeassistant.const import (
+    CONF_DEVICE_ID,
     CONF_USERNAME,
+    CONF_ACCESS_TOKEN,
     CONF_PASSWORD,
 )
 

custom_components/pumpspy_ha/const.py

@@ -4,6 +4,7 @@
 
 
 DOMAIN = "pumpspy_ha"
+MANUFACTURER = "Pumpspy"
 
 BASE_URL = "http://www.pumpspy.com:8081"
 TOKEN_URL = "/oauth/token"

custom_components/pumpspy_ha/entity.py

@@ -0,0 +1,29 @@
+"""Base Entity for Pumpspy."""
+from __future__ import annotations
+
+from homeassistant.helpers.entity import DeviceInfo
+from . import PumpspyCoordinator
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from .const import DOMAIN, MANUFACTURER
+
+
+class PumpspyEntity(CoordinatorEntity[PumpspyCoordinator]):
+    """Defines a base Pumpspy entity."""
+
+    def __init__(self, *, coordinator: PumpspyCoordinator) -> None:
+        """Initialize the entity."""
+        super().__init__(coordinator)
+
+    @property
+    def device_info(self) -> DeviceInfo | None:
+        try:
+            return DeviceInfo(
+                identifiers={(DOMAIN, self.coordinator.data["current"][0]["deviceid"])},
+                name=self.coordinator.data["current"][0]["user_nickname"],
+                manufacturer=MANUFACTURER,
+                model=self.coordinator.data["current"][0]["device_types_name"],
+                hw_version=self.coordinator.data["current"][0]["hardware_rev"],
+                sw_version=self.coordinator.data["current"][0]["firmware_rev"],
+            )
+        except TypeError:
+            return None

custom_components/pumpspy_ha/manifest.json

@@ -2,8 +2,7 @@
   "domain": "pumpspy_ha",
   "name": "Pumpspy-HA",
   "config_flow": true,
-  "documentation": "https://github.com/Crewski/Pumpspy-HA",
-  "issue_tracker": "https://github.com/Crewski/Pumpspy-HA/issues",
+  "documentation": "https://www.home-assistant.io/integrations/pumpspy_ha",
   "requirements": [],
   "ssdp": [],
   "zeroconf": [],

custom_components/pumpspy_ha/pumpspy.py

@@ -31,12 +31,13 @@ def __init__(self, hass: HomeAssistant) -> None:
         """Initialize."""
         self.hass = hass
         self.username = None
+        self.password = None
         self.access_token = None
         self.refresh_token = None
         self.deviceid = None
         self.device_name = None
 
-    async def get_token(self, password: str) -> None:
+    async def get_token(self) -> None:
         """Get bearer token"""
 
         # GET TOKEN
@@ -47,7 +48,7 @@ async def get_token(self, password: str) -> None:
         data = {
             "grant_type": CONF_PASSWORD,
             CONF_USERNAME: self.username,
-            CONF_PASSWORD: password,
+            CONF_PASSWORD: self.password,
         }
         session = aiohttp_client.async_get_clientsession(self.hass)
         async with session.post(
@@ -79,8 +80,8 @@ async def get_uid(self) -> None:  # GET UID
     async def get_locations(self, username: str, password: str):
         """Get the available locations"""
         self.username = username
-
-        await self.get_token(password)
+        self.password = password
+        await self.get_token()
         uid = await self.get_uid()
         session = aiohttp_client.async_get_clientsession(self.hass)
         async with session.get(
@@ -110,6 +111,9 @@ async def fetch_current_data(self):
         ) as resp:
             if resp.status == 200:
                 return await resp.json()
+            elif resp.status == 401:
+                await self.get_token()
+                self.fetch_current_data()
             else:
                 raise CannotConnect
 
@@ -125,6 +129,9 @@ async def fetch_interval_data(self, type: str, interval: str):
         ) as resp:
             if resp.status == 200:
                 return await resp.json()
+            elif resp.status == 401:
+                await self.get_token()
+                self.fetch_current_data()
             else:
                 raise CannotConnect
 
@@ -152,6 +159,7 @@ def get_variables(self):
             CONF_DEVICEID: self.deviceid,
             CONF_USERNAME: self.username,
             CONF_DEVICE_NAME: self.device_name,
+            CONF_PASSWORD: self.password,
         }
 
     def set_variables(
@@ -161,13 +169,15 @@ def set_variables(
         deviceid,
         username: str,
         device_name: str,
+        password: str,
     ):
         """Set all the variables for a new instance"""
         self.username = username
         self.access_token = access_token
         self.refresh_token = refresh_token
         self.deviceid = deviceid
         self.device_name = device_name
+        self.password = password
 
 
 class CannotConnect(HomeAssistantError):