feat: add data update coordinators

- add account and rig update coordinators
- sensors share data

custom_components/nicehash/__init__.py

@@ -4,7 +4,6 @@ Integrates NiceHash with Home Assistant
 For more details about this integration, please refer to
 https://github.com/brianberg/ha-nicehash
 """
-import asyncio
 from datetime import timedelta
 import logging
 import voluptuous as vol
@@ -14,7 +13,8 @@ from homeassistant.const import CONF_DEVICES, CONF_TIMEOUT
 from homeassistant.core import Config, HomeAssistant
 from homeassistant.helpers import discovery
 import homeassistant.helpers.config_validation as cv
-from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.exceptions import PlatformNotReady
 
 from .const import (
     CONF_API_KEY,
@@ -25,7 +25,10 @@ from .const import (
     DOMAIN,
     STARTUP_MESSAGE,
 )
-from .nicehash import NiceHashPrivateClient
+from .nicehash import NiceHashPrivateClient, NiceHashPublicClient
+
+SCAN_INTERVAL_RIGS = timedelta(minutes=1)
+SCAN_INTERVAL_ACCOUNTS = timedelta(minutes=60)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -58,12 +61,84 @@ async def async_setup(hass: HomeAssistant, config: Config):
 
     client = NiceHashPrivateClient(organization_id, api_key, api_secret)
 
-    try:
+    accounts_coordinator = NiceHashAccountsDataUpdateCoordinator(hass, client)
-        await client.get_accounts()
+    rigs_coordinator = NiceHashMiningRigsDataUpdateCoordinator(hass, client)
-        hass.data[DOMAIN]["client"] = client
+
-        hass.data[DOMAIN]["currency"] = currency
+    await accounts_coordinator.async_refresh()
-        await discovery.async_load_platform(hass, "sensor", DOMAIN, {}, config)
+
-        return True
+    if not accounts_coordinator.last_update_success:
-    except Exception as err:
+        _LOGGER.error("Unable to get NiceHash accounts")
-        _LOGGER.error(f"Unable to access NiceHash accounts\n{err}")
+        raise PlatformNotReady
-        return False
+
+    await rigs_coordinator.async_refresh()
+
+    if not rigs_coordinator.last_update_success:
+        _LOGGER.error("Unable to get NiceHash mining rigs")
+        raise PlatformNotReady
+
+    hass.data[DOMAIN]["client"] = client
+    hass.data[DOMAIN]["currency"] = currency
+    hass.data[DOMAIN]["accounts_coordinator"] = accounts_coordinator
+    hass.data[DOMAIN]["rigs_coordinator"] = rigs_coordinator
+
+    await discovery.async_load_platform(hass, "sensor", DOMAIN, {}, config)
+
+    return True
+
+
+class NiceHashAccountsDataUpdateCoordinator(DataUpdateCoordinator):
+    """Manages fetching accounts data from NiceHash API"""
+
+    def __init__(self, hass: HomeAssistant, client: NiceHashPrivateClient):
+        """Initialize"""
+        self.name = f"{DOMAIN}_Accounts_Coordinator"
+        self._client = client
+
+        super().__init__(
+            hass, _LOGGER, name=self.name, update_interval=SCAN_INTERVAL_ACCOUNTS
+        )
+
+    async def _async_update_data(self):
+        """Update accounts data and exchange rates"""
+        try:
+            accounts = await self._client.get_accounts()
+            exchange_rates = await NiceHashPublicClient().get_exchange_rates()
+            rates_dict = dict()
+            for rate in exchange_rates:
+                fromCurrency = rate.get("fromCurrency")
+                toCurrency = rate.get("toCurrency")
+                exchange_rate = float(rate.get("exchangeRate"))
+                rates_dict[f"{fromCurrency}-{toCurrency}"] = exchange_rate
+            return {
+                "accounts": accounts,
+                "exchange_rates": rates_dict,
+            }
+        except Exception as e:
+            raise UpdateFailed(e)
+
+
+class NiceHashMiningRigsDataUpdateCoordinator(DataUpdateCoordinator):
+    """Manages fetching mining rigs data from NiceHash API"""
+
+    def __init__(self, hass: HomeAssistant, client: NiceHashPrivateClient):
+        """Initialize"""
+        self.name = f"{DOMAIN}_Mining_Rigs_Coordinator"
+        self._client = client
+
+        super().__init__(
+            hass, _LOGGER, name=self.name, update_interval=SCAN_INTERVAL_RIGS
+        )
+
+    async def _async_update_data(self):
+        """Update mining rigs data"""
+        try:
+            data = await self._client.get_mining_rigs()
+            mining_rigs = data.get("miningRigs")
+            rigs_dict = dict()
+            for rig in mining_rigs:
+                rig_id = rig.get("rigId")
+                rigs_dict[f"{rig_id}"] = rig
+            data["miningRigs"] = rigs_dict
+            return data
+        except Exception as e:
+            raise UpdateFailed(e)

custom_components/nicehash/sensor.py

@@ -13,6 +13,8 @@ from homeassistant.util import Throttle
 
 from .const import (
     CURRENCY_BTC,
+    CURRENCY_EUR,
+    CURRENCY_USD,
     DEFAULT_NAME,
     DOMAIN,
     ICON_CURRENCY_BTC,
@@ -33,23 +35,35 @@ async def async_setup_platform(
     """Setup NiceHash sensor platform"""
     _LOGGER.debug("Creating new NiceHash sensor components")
 
-    client = hass.data[DOMAIN]["client"]
+    data = hass.data[DOMAIN]
-    currency = hass.data[DOMAIN]["currency"]
+    client = data.get("client")
+    currency = data.get("currency")
+    accounts_coordinator = data.get("accounts_coordinator")
+    rigs_coordinator = data.get("rigs_coordinator")
 
-    # Add account balance sensor
+    # Add account balance sensor(s)
-    async_add_entities(
+    btc_balance_sensor = NiceHashBalanceSensor(
-        [NiceHashBalanceSensor(client, currency, SCAN_INTERVAL_ACCOUNTS)], True
+        accounts_coordinator, client.organization_id, CURRENCY_BTC
     )
+    if currency == CURRENCY_BTC:
+        async_add_entities([btc_balance_sensor], True)
+    else:
+        async_add_entities(
+            [
+                btc_balance_sensor,
+                NiceHashBalanceSensor(
+                    accounts_coordinator, client.organization_id, currency
+                ),
+            ],
+            True,
+        )
 
     # Add mining rig sensors
     rig_data = await client.get_mining_rigs()
     mining_rigs = rig_data["miningRigs"]
     # Add temperature sensors
     async_add_entities(
-        [
+        [NiceHashRigTemperatureSensor(rigs_coordinator, rig) for rig in mining_rigs],
-            NiceHashRigTemperatureSensor(client, rig, SCAN_INTERVAL_RIGS)
-            for rig in mining_rigs
-        ],
         True,
     )
 
@@ -57,82 +71,94 @@ async def async_setup_platform(
 class NiceHashBalanceSensor(Entity):
     """NiceHash Account Balance Sensor"""
 
-    def __init__(self, client, currency, update_frequency):
+    def __init__(self, coordinator, organization_id, currency):
         """Initialize the sensor"""
         _LOGGER.debug(f"Account Balance Sensor: {currency}")
-        self._client = client
+        self.coordinator = coordinator
-        self._public_client = NiceHashPublicClient()
+        self.currency = currency
-        self._currency = currency
+        self._organization_id = organization_id
-        self._state = None
+        self._available = 0.00
-        self._last_update = None
+        self._pending = 0.00
-        self.async_update = Throttle(update_frequency)(self._async_update)
+        self._total_balance = 0.00
+        self._exchange_rate = 0.00
 
     @property
     def name(self):
         """Sensor name"""
-        return f"{DEFAULT_NAME} Account Balance"
+        return f"{DEFAULT_NAME} Account Balance {self.currency}"
 
     @property
     def unique_id(self):
         """Unique entity id"""
-        return f"{self._client.organization_id}:{self._currency}"
+        return f"{self._organization_id}:{self.currency}"
+
+    @property
+    def should_poll(self):
+        """No need to pool, Coordinator notifies entity of updates"""
+        return False
+
+    @property
+    def available(self):
+        """Whether sensor is available"""
+        return self.coordinator.last_update_success
+
+    @property
+    def state(self):
+        """Sensor state"""
+        accounts = self.coordinator.data.get("accounts")
+        total = accounts.get("total")
+        self._pending = float(total.get("pending"))
+        self._available = float(total.get("available"))
+        self._total_balance = float(total.get("totalBalance"))
+        if self.currency == CURRENCY_BTC:
+            return self._available
+        else:
+            exchange_rates = self.coordinator.data.get("exchange_rates")
+            self._exchange_rate = exchange_rates.get(f"{CURRENCY_BTC}-{self.currency}")
+            return round(self._available * self._exchange_rate, 2)
 
     @property
     def icon(self):
         """Sensor icon"""
         return ICON_CURRENCY_BTC
 
-    @property
-    def state(self):
-        """Sensor state"""
-        return self._state
-
     @property
     def unit_of_measurement(self):
         """Sensor unit of measurement"""
-        return self._currency
+        return self.currency
 
     @property
     def device_state_attributes(self):
-        """Sensor device attributes"""
+        """Sensor device state attributes"""
-        return {"last_update": self._last_update}
+        return {
+            "total": self._total_balance,
+            "available": self._available,
+            "pending": self._pending,
+            "exchange_rate": self._exchange_rate,
+        }
 
-    async def _async_update(self):
+    async def async_added_to_hass(self):
+        """Connect to dispatcher listening for entity data notifications"""
+        self.async_on_remove(
+            self.coordinator.async_add_listener(self.async_write_ha_state)
+        )
 
-        try:
+    async def async_update(self):
-            account_data = await self._client.get_accounts()
+        """Update entity"""
-            available = float(account_data["total"]["available"])
+        await self.coordinator.async_request_refresh()
-
-            if self._currency == CURRENCY_BTC:
-                # Account balance is in BTC
-                self._state = available
-            else:
-                # Convert to selected currency via exchange rates
-                exchange_rates = await self._public_client.get_exchange_rates()
-                self._last_update = datetime.today().strftime(FORMAT_DATETIME)
-                for rate in exchange_rates:
-                    isBTC = rate["fromCurrency"] == CURRENCY_BTC
-                    toCurrency = rate["toCurrency"]
-                    if isBTC and toCurrency == self._currency:
-                        rate = float(rate["exchangeRate"])
-                        self._state = round(available * rate, 2)
-        except Exception as err:
-            _LOGGER.error(f"Unable to get account balance\n{err}")
-            pass
 
 
 class NiceHashRigTemperatureSensor(Entity):
     """NichHash Mining Rig Temperature Sensor"""
 
-    def __init__(self, client, rig, update_frequency):
+    def __init__(self, coordinator, rig):
         """Initialize the sensor"""
-        self._client = client
+        self.coordinator = coordinator
         self._rig_id = rig["rigId"]
         self._name = rig["name"]
+        self._temps = []
+        self._num_devices = 0
         _LOGGER.debug(f"Mining Rig Temperature Sensor: {self._name} ({self._rig_id})")
-        self._state = None
-        self._last_update = None
-        self.async_update = Throttle(update_frequency)(self._async_update)
 
     @property
     def name(self):
@@ -144,10 +170,46 @@ class NiceHashRigTemperatureSensor(Entity):
         """Unique entity id"""
         return self._rig_id
 
+    @property
+    def should_poll(self):
+        """No need to pool, Coordinator notifies entity of updates"""
+        return False
+
+    @property
+    def available(self):
+        """Whether sensor is available"""
+        return self.coordinator.last_update_success
+
     @property
     def state(self):
         """Sensor state"""
-        return self._state
+        mining_rigs = self.coordinator.data.get("miningRigs")
+        try:
+            rig_data = mining_rigs.get(self._rig_id)
+            devices = rig_data.get("devices")
+            highest_temp = 0
+            num_devices = len(devices)
+            self._num_devices = num_devices
+
+            if num_devices > 0:
+                _LOGGER.debug(f"{self._name}: Found {num_devices} devices")
+                self._temps = []
+                for device in devices:
+                    temp = int(device.get("temperature"))
+                    self._temps.append(temp)
+                    if temp < 0:
+                        # Ignore inactive devices
+                        continue
+                    if temp > highest_temp:
+                        highest_temp = temp
+                    return highest_temp
+            else:
+                _LOGGER.debug(f"{self._name}: No devices found")
+                self._num_devices = 0
+                return 0
+        except Exception as e:
+            _LOGGER.error(f"Unable to get mining rig {self._rig_id}\n{e}")
+            return 0
 
     @property
     def icon(self):
@@ -161,32 +223,18 @@ class NiceHashRigTemperatureSensor(Entity):
 
     @property
     def device_state_attributes(self):
-        """Sensore device state attributes"""
+        """Sensor device state attributes"""
         return {
-            "last_update": self._last_update,
+            "temperatures": self._temps,
+            "num_devices": self._num_devices,
         }
 
-    async def _async_update(self):
+    async def async_added_to_hass(self):
-        try:
+        """Connect to dispatcher listening for entity data notifications"""
-            data = await self._client.get_mining_rig(self._rig_id)
+        self.async_on_remove(
-            self._last_update = datetime.today().strftime(FORMAT_DATETIME)
+            self.coordinator.async_add_listener(self.async_write_ha_state)
-            devices = data["devices"]
+        )
-            highest_temp = 0
 
-            if len(devices) > 0:
+    async def async_update(self):
-                _LOGGER.debug(f"{self._name}: Found {len(devices)} devices")
+        """Update entity"""
-                for device in devices:
+        await self.coordinator.async_request_refresh()
-                    temp = int(device["temperature"])
-                    if temp < 0:
-                        # Ignore inactive devices
-                        continue
-                    if temp > highest_temp:
-                        highest_temp = temp
-                    self._state = highest_temp
-            else:
-                _LOGGER.debug(f"{self._name}: No devices found")
-                self._state = None
-
-        except Exception as err:
-            _LOGGER.error(f"Unable to get mining rig {self._rig_id}\n{err}")
-            pass