Add account sensors to electric kiwi integration (#97681)

* add account sensors

* tidy up same issues as other sensors

* add unit tests for sensors

edit and remove comments

assert state and remove HOP sensor types since they aren't being used

* try and fix tests

* add frozen true

* Update tests/components/electric_kiwi/test_sensor.py

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* return proper native types

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* tidy up attr unique id

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* add entities once and use native values properly

* Improve conftest

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* tidy tests/components/electric_kiwi/test_sensor.py

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* add assert to component_setup

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* add extra parameters to test

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* Update tests/components/electric_kiwi/test_sensor.py

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* Update tests/components/electric_kiwi/test_sensor.py

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* change coordinator name

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

* tidy up sensor translation names

* Apply suggestions from code review

---------

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

homeassistant/components/electric_kiwi/__init__.py

@@ -12,8 +12,11 @@ from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow
 
 from . import api
-from .const import DOMAIN
+from .const import ACCOUNT_COORDINATOR, DOMAIN, HOP_COORDINATOR
-from .coordinator import ElectricKiwiHOPDataCoordinator
+from .coordinator import (
+    ElectricKiwiAccountDataCoordinator,
+    ElectricKiwiHOPDataCoordinator,
+)
 
 PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.SELECT]
 
@@ -41,14 +44,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         api.AsyncConfigEntryAuth(aiohttp_client.async_get_clientsession(hass), session)
     )
     hop_coordinator = ElectricKiwiHOPDataCoordinator(hass, ek_api)
+    account_coordinator = ElectricKiwiAccountDataCoordinator(hass, ek_api)
 
     try:
         await ek_api.set_active_session()
         await hop_coordinator.async_config_entry_first_refresh()
+        await account_coordinator.async_config_entry_first_refresh()
     except ApiException as err:
         raise ConfigEntryNotReady from err
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = hop_coordinator
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
+        HOP_COORDINATOR: hop_coordinator,
+        ACCOUNT_COORDINATOR: account_coordinator,
+    }
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 

homeassistant/components/electric_kiwi/const.py

@@ -9,3 +9,6 @@ OAUTH2_TOKEN = "https://welcome.electrickiwi.co.nz/oauth/token"
 API_BASE_URL = "https://api.electrickiwi.co.nz"
 
 SCOPE_VALUES = "read_connection_detail read_billing_frequency read_account_running_balance read_consumption_summary read_consumption_averages read_hop_intervals_config read_hop_connection save_hop_connection read_session"
+
+HOP_COORDINATOR = "hop_coordinator"
+ACCOUNT_COORDINATOR = "account_coordinator"

homeassistant/components/electric_kiwi/coordinator.py

@@ -6,7 +6,7 @@ import logging
 
 from electrickiwi_api import ElectricKiwiApi
 from electrickiwi_api.exceptions import ApiException, AuthException
-from electrickiwi_api.model import Hop, HopIntervals
+from electrickiwi_api.model import AccountBalance, Hop, HopIntervals
 
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
@@ -14,11 +14,38 @@ from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, Upda
 
 _LOGGER = logging.getLogger(__name__)
 
+ACCOUNT_SCAN_INTERVAL = timedelta(hours=6)
 HOP_SCAN_INTERVAL = timedelta(minutes=20)
 
 
+class ElectricKiwiAccountDataCoordinator(DataUpdateCoordinator):
+    """ElectricKiwi Account Data object."""
+
+    def __init__(self, hass: HomeAssistant, ek_api: ElectricKiwiApi) -> None:
+        """Initialize ElectricKiwiAccountDataCoordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name="Electric Kiwi Account Data",
+            update_interval=ACCOUNT_SCAN_INTERVAL,
+        )
+        self._ek_api = ek_api
+
+    async def _async_update_data(self) -> AccountBalance:
+        """Fetch data from Account balance API endpoint."""
+        try:
+            async with asyncio.timeout(60):
+                return await self._ek_api.get_account_balance()
+        except AuthException as auth_err:
+            raise ConfigEntryAuthFailed from auth_err
+        except ApiException as api_err:
+            raise UpdateFailed(
+                f"Error communicating with EK API: {api_err}"
+            ) from api_err
+
+
 class ElectricKiwiHOPDataCoordinator(DataUpdateCoordinator[Hop]):
-    """ElectricKiwi Data object."""
+    """ElectricKiwi HOP Data object."""
 
     def __init__(self, hass: HomeAssistant, ek_api: ElectricKiwiApi) -> None:
         """Initialize ElectricKiwiAccountDataCoordinator."""

homeassistant/components/electric_kiwi/select.py

@@ -10,7 +10,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import ATTRIBUTION, DOMAIN
+from .const import ATTRIBUTION, DOMAIN, HOP_COORDINATOR
 from .coordinator import ElectricKiwiHOPDataCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -19,7 +19,7 @@ ATTR_EK_HOP_SELECT = "hop_select"
 HOP_SELECT = SelectEntityDescription(
     entity_category=EntityCategory.CONFIG,
     key=ATTR_EK_HOP_SELECT,
-    translation_key="hopselector",
+    translation_key="hop_selector",
 )
 
 
@@ -27,7 +27,9 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Electric Kiwi select setup."""
-    hop_coordinator: ElectricKiwiHOPDataCoordinator = hass.data[DOMAIN][entry.entry_id]
+    hop_coordinator: ElectricKiwiHOPDataCoordinator = hass.data[DOMAIN][entry.entry_id][
+        HOP_COORDINATOR
+    ]
 
     _LOGGER.debug("Setting up select entity")
     async_add_entities([ElectricKiwiSelectHOPEntity(hop_coordinator, HOP_SELECT)])

homeassistant/components/electric_kiwi/sensor.py

@@ -4,28 +4,89 @@ from __future__ import annotations
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime, timedelta
-import logging
 
-from electrickiwi_api.model import Hop
+from electrickiwi_api.model import AccountBalance, Hop
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
     SensorEntityDescription,
+    SensorStateClass,
 )
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CURRENCY_DOLLAR, PERCENTAGE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.util import dt as dt_util
 
-from .const import ATTRIBUTION, DOMAIN
+from .const import ACCOUNT_COORDINATOR, ATTRIBUTION, DOMAIN, HOP_COORDINATOR
-from .coordinator import ElectricKiwiHOPDataCoordinator
+from .coordinator import (
+    ElectricKiwiAccountDataCoordinator,
+    ElectricKiwiHOPDataCoordinator,
+)
 
-_LOGGER = logging.getLogger(DOMAIN)
+ATTR_EK_HOP_START = "hop_power_start"
+ATTR_EK_HOP_END = "hop_power_end"
+ATTR_TOTAL_RUNNING_BALANCE = "total_running_balance"
+ATTR_TOTAL_CURRENT_BALANCE = "total_account_balance"
+ATTR_NEXT_BILLING_DATE = "next_billing_date"
+ATTR_HOP_PERCENTAGE = "hop_percentage"
 
-ATTR_EK_HOP_START = "hop_sensor_start"
+
-ATTR_EK_HOP_END = "hop_sensor_end"
+@dataclass(frozen=True)
+class ElectricKiwiAccountRequiredKeysMixin:
+    """Mixin for required keys."""
+
+    value_func: Callable[[AccountBalance], float | datetime]
+
+
+@dataclass(frozen=True)
+class ElectricKiwiAccountSensorEntityDescription(
+    SensorEntityDescription, ElectricKiwiAccountRequiredKeysMixin
+):
+    """Describes Electric Kiwi sensor entity."""
+
+
+ACCOUNT_SENSOR_TYPES: tuple[ElectricKiwiAccountSensorEntityDescription, ...] = (
+    ElectricKiwiAccountSensorEntityDescription(
+        key=ATTR_TOTAL_RUNNING_BALANCE,
+        translation_key="total_running_balance",
+        icon="mdi:currency-usd",
+        device_class=SensorDeviceClass.MONETARY,
+        state_class=SensorStateClass.TOTAL,
+        native_unit_of_measurement=CURRENCY_DOLLAR,
+        value_func=lambda account_balance: float(account_balance.total_running_balance),
+    ),
+    ElectricKiwiAccountSensorEntityDescription(
+        key=ATTR_TOTAL_CURRENT_BALANCE,
+        translation_key="total_current_balance",
+        icon="mdi:currency-usd",
+        device_class=SensorDeviceClass.MONETARY,
+        state_class=SensorStateClass.TOTAL,
+        native_unit_of_measurement=CURRENCY_DOLLAR,
+        value_func=lambda account_balance: float(account_balance.total_account_balance),
+    ),
+    ElectricKiwiAccountSensorEntityDescription(
+        key=ATTR_NEXT_BILLING_DATE,
+        translation_key="next_billing_date",
+        icon="mdi:calendar",
+        device_class=SensorDeviceClass.DATE,
+        value_func=lambda account_balance: datetime.strptime(
+            account_balance.next_billing_date, "%Y-%m-%d"
+        ),
+    ),
+    ElectricKiwiAccountSensorEntityDescription(
+        key=ATTR_HOP_PERCENTAGE,
+        translation_key="hop_power_savings",
+        icon="mdi:percent",
+        native_unit_of_measurement=PERCENTAGE,
+        state_class=SensorStateClass.MEASUREMENT,
+        value_func=lambda account_balance: float(
+            account_balance.connections[0].hop_percentage
+        ),
+    ),
+)
 
 
 @dataclass(frozen=True)
@@ -65,13 +126,13 @@ def _check_and_move_time(hop: Hop, time: str) -> datetime:
 HOP_SENSOR_TYPES: tuple[ElectricKiwiHOPSensorEntityDescription, ...] = (
     ElectricKiwiHOPSensorEntityDescription(
         key=ATTR_EK_HOP_START,
-        translation_key="hopfreepowerstart",
+        translation_key="hop_free_power_start",
         device_class=SensorDeviceClass.TIMESTAMP,
         value_func=lambda hop: _check_and_move_time(hop, hop.start.start_time),
     ),
     ElectricKiwiHOPSensorEntityDescription(
         key=ATTR_EK_HOP_END,
-        translation_key="hopfreepowerend",
+        translation_key="hop_free_power_end",
         device_class=SensorDeviceClass.TIMESTAMP,
         value_func=lambda hop: _check_and_move_time(hop, hop.end.end_time),
     ),
@@ -81,13 +142,58 @@ HOP_SENSOR_TYPES: tuple[ElectricKiwiHOPSensorEntityDescription, ...] = (
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
-    """Electric Kiwi Sensor Setup."""
+    """Electric Kiwi Sensors Setup."""
-    hop_coordinator: ElectricKiwiHOPDataCoordinator = hass.data[DOMAIN][entry.entry_id]
+    account_coordinator: ElectricKiwiAccountDataCoordinator = hass.data[DOMAIN][
-    hop_entities = [
+        entry.entry_id
+    ][ACCOUNT_COORDINATOR]
+
+    entities: list[SensorEntity] = [
+        ElectricKiwiAccountEntity(
+            account_coordinator,
+            description,
+        )
+        for description in ACCOUNT_SENSOR_TYPES
+    ]
+
+    hop_coordinator: ElectricKiwiHOPDataCoordinator = hass.data[DOMAIN][entry.entry_id][
+        HOP_COORDINATOR
+    ]
+    entities.extend(
+        [
             ElectricKiwiHOPEntity(hop_coordinator, description)
             for description in HOP_SENSOR_TYPES
         ]
-    async_add_entities(hop_entities)
+    )
+    async_add_entities(entities)
+
+
+class ElectricKiwiAccountEntity(
+    CoordinatorEntity[ElectricKiwiAccountDataCoordinator], SensorEntity
+):
+    """Entity object for Electric Kiwi sensor."""
+
+    entity_description: ElectricKiwiAccountSensorEntityDescription
+    _attr_has_entity_name = True
+    _attr_attribution = ATTRIBUTION
+
+    def __init__(
+        self,
+        coordinator: ElectricKiwiAccountDataCoordinator,
+        description: ElectricKiwiAccountSensorEntityDescription,
+    ) -> None:
+        """Entity object for Electric Kiwi sensor."""
+        super().__init__(coordinator)
+
+        self._attr_unique_id = (
+            f"{coordinator._ek_api.customer_number}"
+            f"_{coordinator._ek_api.connection_id}_{description.key}"
+        )
+        self.entity_description = description
+
+    @property
+    def native_value(self) -> float | datetime:
+        """Return the state of the sensor."""
+        return self.entity_description.value_func(self.coordinator.data)
 
 
 class ElectricKiwiHOPEntity(

homeassistant/components/electric_kiwi/strings.json

@@ -28,9 +28,25 @@
   },
   "entity": {
     "sensor": {
-      "hopfreepowerstart": { "name": "Hour of free power start" },
+      "hop_free_power_start": {
-      "hopfreepowerend": { "name": "Hour of free power end" }
+        "name": "Hour of free power start"
       },
-    "select": { "hopselector": { "name": "Hour of free power" } }
+      "hop_free_power_end": {
+        "name": "Hour of free power end"
+      },
+      "total_running_balance": {
+        "name": "Total running balance"
+      },
+      "total_current_balance": {
+        "name": "Total current balance"
+      },
+      "next_billing_date": {
+        "name": "Next billing date"
+      },
+      "hop_power_savings": {
+        "name": "Hour of power savings"
+      }
+    },
+    "select": { "hop_selector": { "name": "Hour of free power" } }
   }
 }

tests/components/electric_kiwi/conftest.py

@@ -6,7 +6,7 @@ from time import time
 from unittest.mock import AsyncMock, patch
 import zoneinfo
 
-from electrickiwi_api.model import Hop, HopIntervals
+from electrickiwi_api.model import AccountBalance, Hop, HopIntervals
 import pytest
 
 from homeassistant.components.application_credentials import (
@@ -43,14 +43,18 @@ def component_setup(
 
     async def _setup_func() -> bool:
         assert await async_setup_component(hass, "application_credentials", {})
+        await hass.async_block_till_done()
         await async_import_client_credential(
             hass,
             DOMAIN,
             ClientCredential(CLIENT_ID, CLIENT_SECRET),
             DOMAIN,
         )
+        await hass.async_block_till_done()
         config_entry.add_to_hass(hass)
-        return await hass.config_entries.async_setup(config_entry.entry_id)
+        result = await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+        return result
 
     return _setup_func
 
@@ -113,4 +117,9 @@ def ek_api() -> YieldFixture:
         mock_ek_api.return_value.get_hop.return_value = Hop.from_dict(
             load_json_value_fixture("get_hop.json", DOMAIN)
         )
+        mock_ek_api.return_value.get_account_balance.return_value = (
+            AccountBalance.from_dict(
+                load_json_value_fixture("account_balance.json", DOMAIN)
+            )
+        )
         yield mock_ek_api

tests/components/electric_kiwi/fixtures/account_balance.json

@@ -0,0 +1,28 @@
+{
+  "data": {
+    "connections": [
+      {
+        "hop_percentage": "3.5",
+        "id": 3,
+        "running_balance": "184.09",
+        "start_date": "2020-10-04",
+        "unbilled_days": 15
+      }
+    ],
+    "last_billed_amount": "-66.31",
+    "last_billed_date": "2020-10-03",
+    "next_billing_date": "2020-11-03",
+    "is_prepay": "N",
+    "summary": {
+      "credits": "0.0",
+      "electricity_used": "184.09",
+      "other_charges": "0.00",
+      "payments": "-220.0"
+    },
+    "total_account_balance": "-102.22",
+    "total_billing_days": 30,
+    "total_running_balance": "184.09",
+    "type": "account_running_balance"
+  },
+  "status": 1
+}

tests/components/electric_kiwi/test_sensor.py

@@ -9,7 +9,11 @@ import pytest
 
 from homeassistant.components.electric_kiwi.const import ATTRIBUTION
 from homeassistant.components.electric_kiwi.sensor import _check_and_move_time
-from homeassistant.components.sensor import SensorDeviceClass
+from homeassistant.components.sensor import (
+    ATTR_STATE_CLASS,
+    SensorDeviceClass,
+    SensorStateClass,
+)
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import ATTR_ATTRIBUTION, ATTR_DEVICE_CLASS
 from homeassistant.core import HomeAssistant
@@ -65,6 +69,58 @@ async def test_hop_sensors(
     assert state.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.TIMESTAMP
 
 
+@pytest.mark.parametrize(
+    ("sensor", "sensor_state", "device_class", "state_class"),
+    [
+        (
+            "sensor.total_running_balance",
+            "184.09",
+            SensorDeviceClass.MONETARY,
+            SensorStateClass.TOTAL,
+        ),
+        (
+            "sensor.total_current_balance",
+            "-102.22",
+            SensorDeviceClass.MONETARY,
+            SensorStateClass.TOTAL,
+        ),
+        (
+            "sensor.next_billing_date",
+            "2020-11-03T00:00:00",
+            SensorDeviceClass.DATE,
+            None,
+        ),
+        ("sensor.hour_of_power_savings", "3.5", None, SensorStateClass.MEASUREMENT),
+    ],
+)
+async def test_account_sensors(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    ek_api: YieldFixture,
+    ek_auth: YieldFixture,
+    entity_registry: EntityRegistry,
+    component_setup: ComponentSetup,
+    sensor: str,
+    sensor_state: str,
+    device_class: str,
+    state_class: str,
+) -> None:
+    """Test Account sensors for the Electric Kiwi integration."""
+
+    assert await component_setup()
+    assert config_entry.state is ConfigEntryState.LOADED
+
+    entity = entity_registry.async_get(sensor)
+    assert entity
+
+    state = hass.states.get(sensor)
+    assert state
+    assert state.state == sensor_state
+    assert state.attributes.get(ATTR_ATTRIBUTION) == ATTRIBUTION
+    assert state.attributes.get(ATTR_DEVICE_CLASS) == device_class
+    assert state.attributes.get(ATTR_STATE_CLASS) == state_class
+
+
 async def test_check_and_move_time(ek_api: AsyncMock) -> None:
     """Test correct time is returned depending on time of day."""
     hop = await ek_api(Mock()).get_hop()