Add energy usage sensor and tweak switch update logic (#3)

dan-r · web-flow · commit a5bd807811a8 · 2023-12-27T19:48:59.000Z
* Tweak switch update logic

* Added accumulative energy usage sensor

* Properly implement unique_id

* Changed energy usage sensor to kWh
diff --git a/custom_components/ohme/__init__.py b/custom_components/ohme/__init__.py
@@ -1,7 +1,7 @@
 from homeassistant import core
 from .const import *
 from .client import OhmeApiClient
-from .coordinator import OhmeUpdateCoordinator
+from .coordinator import OhmeUpdateCoordinator, OhmeStatisticsUpdateCoordinator
 
 
 async def async_setup(hass: core.HomeAssistant, config: dict) -> bool:
@@ -30,7 +30,14 @@ async def async_setup_entry(hass, entry):
         return False
 
     await async_setup_dependencies(hass, config)
+    
+    hass.data[DOMAIN][DATA_COORDINATOR] = OhmeUpdateCoordinator(hass=hass)
+    await hass.data[DOMAIN][DATA_COORDINATOR].async_config_entry_first_refresh()
 
+    hass.data[DOMAIN][DATA_STATISTICS_COORDINATOR] = OhmeStatisticsUpdateCoordinator(hass=hass)
+    await hass.data[DOMAIN][DATA_STATISTICS_COORDINATOR].async_config_entry_first_refresh()
+
+    # Create tasks for each entity type
     hass.async_create_task(
         hass.config_entries.async_forward_entry_setup(entry, "sensor")
     )
@@ -41,7 +48,9 @@ async def async_setup_entry(hass, entry):
         hass.config_entries.async_forward_entry_setup(entry, "switch")
     )
 
-    hass.data[DOMAIN][DATA_COORDINATOR] = OhmeUpdateCoordinator(hass=hass)
-    await hass.data[DOMAIN][DATA_COORDINATOR].async_config_entry_first_refresh()
-
     return True
+
+async def async_unload_entry(hass, entry):
+    """Unload a config entry."""
+
+    return await hass.config_entries.async_unload_platforms(entry, ['binary_sensor', 'sensor', 'switch'])
diff --git a/custom_components/ohme/binary_sensor.py b/custom_components/ohme/binary_sensor.py
@@ -18,10 +18,11 @@ async def async_setup_entry(
     async_add_entities,
 ):
     """Setup sensors and configure coordinator."""
+    client = hass.data[DOMAIN][DATA_CLIENT]
     coordinator = hass.data[DOMAIN][DATA_COORDINATOR]
 
-    sensors = [ConnectedSensor(coordinator, hass),
-               ChargingSensor(coordinator, hass)]
+    sensors = [ConnectedSensor(coordinator, hass, client),
+               ChargingSensor(coordinator, hass, client)]
 
     async_add_entities(sensors, update_before_add=True)
 
@@ -37,12 +38,14 @@ class ConnectedSensor(
     def __init__(
             self,
             coordinator: OhmeUpdateCoordinator,
-            hass: HomeAssistant):
+            hass: HomeAssistant,
+            client):
         super().__init__(coordinator=coordinator)
 
         self._attributes = {}
         self._last_updated = None
         self._state = False
+        self._client = client
 
         self.entity_id = generate_entity_id(
             "binary_sensor.{}", "ohme_car_connected", hass=hass)
@@ -58,7 +61,7 @@ def icon(self):
     @property
     def unique_id(self) -> str:
         """Return the unique ID of the sensor."""
-        return self.entity_id
+        return self._client.get_unique_id("car_connected")
 
     @property
     def is_on(self) -> bool:
@@ -81,12 +84,14 @@ class ChargingSensor(
     def __init__(
             self,
             coordinator: OhmeUpdateCoordinator,
-            hass: HomeAssistant):
+            hass: HomeAssistant,
+            client):
         super().__init__(coordinator=coordinator)
 
         self._attributes = {}
         self._last_updated = None
         self._state = False
+        self._client = client
 
         self.entity_id = generate_entity_id(
             "binary_sensor.{}", "ohme_car_charging", hass=hass)
@@ -102,12 +107,13 @@ def icon(self):
     @property
     def unique_id(self) -> str:
         """Return the unique ID of the sensor."""
-        return self.entity_id
+        return self._client.get_unique_id("ohme_car_charging")
 
     @property
     def is_on(self) -> bool:
         if self.coordinator.data and self.coordinator.data["power"]:
-            self._state = self.coordinator.data["power"]["amp"] > 0
+            # Assume the car is actively charging if drawing over 0 watts
+            self._state = self.coordinator.data["power"]["watt"] > 0
         else:
             self._state = False
 
diff --git a/custom_components/ohme/client/__init__.py b/custom_components/ohme/client/__init__.py
@@ -2,6 +2,7 @@
 import asyncio
 import logging
 import json
+from datetime import datetime, timedelta
 from homeassistant.helpers.entity import DeviceInfo
 from ..const import DOMAIN
 
@@ -20,6 +21,7 @@ def __init__(self, email, password):
 
         self._device_info = None
         self._token = None
+        self._user_id = ""
         self._serial = ""
         self._session = aiohttp.ClientSession()
 
@@ -74,6 +76,21 @@ async def _put_request(self, url, data=None, is_retry=False):
 
             return True
 
+    async def _get_request(self, url, is_retry=False):
+        """Try to make a GET request
+           If we get a non 200 response, refresh auth token and try again"""
+        async with self._session.get(
+            url,
+            headers={"Authorization": "Firebase %s" % self._token}
+        ) as resp:
+            if resp.status != 200 and not is_retry:
+                await self.async_refresh_session()
+                return await self._get_request(url, is_retry=True)
+            elif resp.status != 200:
+                return False
+
+            return await resp.json()
+
     async def async_pause_charge(self):
         """Pause an ongoing charge"""
         result = await self._post_request(f"https://api.ohme.io/v1/chargeSessions/{self._serial}/stop", skip_json=True)
@@ -98,46 +115,57 @@ async def async_stop_max_charge(self):
     async def async_get_charge_sessions(self, is_retry=False):
         """Try to fetch charge sessions endpoint.
            If we get a non 200 response, refresh auth token and try again"""
-        async with self._session.get(
-            'https://api.ohme.io/v1/chargeSessions',
-            headers={"Authorization": "Firebase %s" % self._token}
-        ) as resp:
-
-            if resp.status != 200 and not is_retry:
-                await self.async_refresh_session()
-                return await self.async_get_charge_sessions(True)
-            elif resp.status != 200:
-                return False
+        resp = await self._get_request('https://api.ohme.io/v1/chargeSessions')
+        
+        if not resp:
+            return False
 
-            resp_json = await resp.json()
-            return resp_json[0]
+        return resp[0]
 
     async def async_update_device_info(self, is_retry=False):
         """Update _device_info with our charger model."""
-        async with self._session.get(
-            'https://api.ohme.io/v1/users/me/account',
-            headers={"Authorization": "Firebase %s" % self._token}
-        ) as resp:
+        resp = await self._get_request('https://api.ohme.io/v1/users/me/account')
 
-            if resp.status != 200 and not is_retry:
-                await self.async_refresh_session()
-                return await self.async_get_device_info(True)
-            elif resp.status != 200:
-                return False
+        if not resp:
+            return False
 
-            resp_json = await resp.json()
-            device = resp_json['chargeDevices'][0]
-
-            info = DeviceInfo(
-                identifiers={(DOMAIN, "ohme_charger")},
-                name=device['modelTypeDisplayName'],
-                manufacturer="Ohme",
-                model=device['modelTypeDisplayName'].replace("Ohme ", ""),
-                sw_version=device['firmwareVersionLabel'],
-                serial_number=device['id']
-            )
-            self._serial = device['id']
-            self._device_info = info
+        device = resp['chargeDevices'][0]
+
+        info = DeviceInfo(
+            identifiers={(DOMAIN, "ohme_charger")},
+            name=device['modelTypeDisplayName'],
+            manufacturer="Ohme",
+            model=device['modelTypeDisplayName'].replace("Ohme ", ""),
+            sw_version=device['firmwareVersionLabel'],
+            serial_number=device['id']
+        )
+            
+        self._user_id = resp['user']['id']
+        self._serial = device['id']
+        self._device_info = info
+
+        return True
+
+    def _last_second_of_month_timestamp(self):
+        """Get the last second of this month."""
+        dt = datetime.today()
+        dt = dt.replace(day=1) + timedelta(days=32)
+        dt = dt.replace(day=1, hour=0, minute=0, second=0, microsecond=0) - timedelta(seconds=1)
+        return int(dt.timestamp()*1e3)
+
+    async def async_get_charge_statistics(self):
+        """Get charge statistics. Currently this is just for all time (well, Jan 2019)."""
+        end_ts = self._last_second_of_month_timestamp()
+        resp = await self._get_request(f"https://api.ohme.io/v1/chargeSessions/summary/users/{self._user_id}?&startTs=1546300800000&endTs={end_ts}&granularity=MONTH")
+
+        if not resp:
+            return False
+
+        return resp['totalStats']
 
     def get_device_info(self):
         return self._device_info
+
+    def get_unique_id(self, name):
+        return f"ohme_{self._serial}_{name}"
+
diff --git a/custom_components/ohme/const.py b/custom_components/ohme/const.py
@@ -3,3 +3,4 @@
 DOMAIN = "ohme"
 DATA_CLIENT = "client"
 DATA_COORDINATOR = "coordinator"
+DATA_STATISTICS_COORDINATOR = "statistics_coordinator"
diff --git a/custom_components/ohme/coordinator.py b/custom_components/ohme/coordinator.py
@@ -31,3 +31,26 @@ async def _async_update_data(self):
 
         except BaseException:
             raise UpdateFailed("Error communicating with API")
+
+
+class OhmeStatisticsUpdateCoordinator(DataUpdateCoordinator):
+    """Coordinator to update statistics from API periodically.
+       (But less so than OhmeUpdateCoordinator)"""
+
+    def __init__(self, hass):
+        """Initialise coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name="Ohme Charger Statistics",
+            update_interval=timedelta(minutes=30),
+        )
+        self._client = hass.data[DOMAIN][DATA_CLIENT]
+
+    async def _async_update_data(self):
+        """Fetch data from API endpoint."""
+        try:
+            return await self._client.async_get_charge_statistics()
+
+        except BaseException:
+            raise UpdateFailed("Error communicating with API")
diff --git a/custom_components/ohme/sensor.py b/custom_components/ohme/sensor.py
@@ -6,11 +6,11 @@
     SensorEntity
 )
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
-from homeassistant.const import UnitOfPower
+from homeassistant.const import UnitOfPower, UnitOfEnergy
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import generate_entity_id
-from .const import DOMAIN, DATA_CLIENT, DATA_COORDINATOR
-from .coordinator import OhmeUpdateCoordinator
+from .const import DOMAIN, DATA_CLIENT, DATA_COORDINATOR, DATA_STATISTICS_COORDINATOR
+from .coordinator import OhmeUpdateCoordinator, OhmeStatisticsUpdateCoordinator
 
 
 async def async_setup_entry(
@@ -19,9 +19,11 @@ async def async_setup_entry(
     async_add_entities
 ):
     """Setup sensors and configure coordinator."""
+    client = hass.data[DOMAIN][DATA_CLIENT]
     coordinator = hass.data[DOMAIN][DATA_COORDINATOR]
+    stats_coordinator = hass.data[DOMAIN][DATA_STATISTICS_COORDINATOR]
 
-    sensors = [PowerDrawSensor(coordinator, hass)]
+    sensors = [PowerDrawSensor(coordinator, hass, client), EnergyUsageSensor(stats_coordinator, hass, client)]
 
     async_add_entities(sensors, update_before_add=True)
 
@@ -35,23 +37,24 @@ class PowerDrawSensor(CoordinatorEntity[OhmeUpdateCoordinator], SensorEntity):
     def __init__(
             self,
             coordinator: OhmeUpdateCoordinator,
-            hass: HomeAssistant):
+            hass: HomeAssistant,
+            client):
         super().__init__(coordinator=coordinator)
 
         self._state = None
         self._attributes = {}
         self._last_updated = None
+        self._client = client
 
         self.entity_id = generate_entity_id(
             "sensor.{}", "ohme_power_draw", hass=hass)
 
-        self._attr_device_info = hass.data[DOMAIN][DATA_CLIENT].get_device_info(
-        )
+        self._attr_device_info = hass.data[DOMAIN][DATA_CLIENT].get_device_info()
 
     @property
     def unique_id(self) -> str:
         """Return the unique ID of the sensor."""
-        return self.entity_id
+        return self._client.get_unique_id("power_draw")
 
     @property
     def icon(self):
@@ -64,3 +67,45 @@ def native_value(self):
         if self.coordinator.data and self.coordinator.data['power']:
             return self.coordinator.data['power']['watt']
         return 0
+
+
+class EnergyUsageSensor(CoordinatorEntity[OhmeStatisticsUpdateCoordinator], SensorEntity):
+    """Sensor for total energy usage."""
+    _attr_name = "Ohme Accumulative Energy Usage"
+    _attr_native_unit_of_measurement = UnitOfEnergy.KILO_WATT_HOUR
+    _attr_device_class = SensorDeviceClass.ENERGY
+
+    def __init__(
+            self,
+            coordinator: OhmeUpdateCoordinator,
+            hass: HomeAssistant,
+            client):
+        super().__init__(coordinator=coordinator)
+
+        self._state = None
+        self._attributes = {}
+        self._last_updated = None
+        self._client = client
+
+        self.entity_id = generate_entity_id(
+            "sensor.{}", "ohme_accumulative_energy", hass=hass)
+
+        self._attr_device_info = hass.data[DOMAIN][DATA_CLIENT].get_device_info()
+
+    @property
+    def unique_id(self) -> str:
+        """Return the unique ID of the sensor."""
+        return self._client.get_unique_id("accumulative_energy")
+
+    @property
+    def icon(self):
+        """Icon of the sensor."""
+        return "mdi:lightning-bolt"
+
+    @property
+    def native_value(self):
+        """Get value from data returned from API by coordinator"""
+        if self.coordinator.data and self.coordinator.data['energyChargedTotalWh']:
+            return self.coordinator.data['energyChargedTotalWh'] / 1000
+
+        return None
diff --git a/custom_components/ohme/switch.py b/custom_components/ohme/switch.py
