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>
2024-01-14 06:12:40 +13:00 · 2024-01-14 06:12:40 +13:00 · b1a246b817
parent 8395d84bbb
commit b1a246b817
9 changed files with 284 additions and 29 deletions
--- a/homeassistant/components/electric_kiwi/init.py
+++ b/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 .coordinator import ElectricKiwiHOPDataCoordinator
+from .const import ACCOUNT_COORDINATOR, DOMAIN, HOP_COORDINATOR
+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)

--- a/homeassistant/components/electric_kiwi/const.py
+++ b/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"
--- a/homeassistant/components/electric_kiwi/coordinator.py
+++ b/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."""
--- a/homeassistant/components/electric_kiwi/select.py
+++ b/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)])
--- a/homeassistant/components/electric_kiwi/sensor.py
+++ b/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 .coordinator import ElectricKiwiHOPDataCoordinator
+from .const import ACCOUNT_COORDINATOR, ATTRIBUTION, DOMAIN, HOP_COORDINATOR
+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."""
-    hop_coordinator: ElectricKiwiHOPDataCoordinator = hass.data[DOMAIN][entry.entry_id]
-    hop_entities = [
-        ElectricKiwiHOPEntity(hop_coordinator, description)
-        for description in HOP_SENSOR_TYPES
+    """Electric Kiwi Sensors Setup."""
+    account_coordinator: ElectricKiwiAccountDataCoordinator = hass.data[DOMAIN][
+        entry.entry_id
+    ][ACCOUNT_COORDINATOR]
+
+    entities: list[SensorEntity] = [
+        ElectricKiwiAccountEntity(
+            account_coordinator,
+            description,
+        )
+        for description in ACCOUNT_SENSOR_TYPES
    ]
-    async_add_entities(hop_entities)
+
+    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(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(
--- a/homeassistant/components/electric_kiwi/strings.json
+++ b/homeassistant/components/electric_kiwi/strings.json
@ -28,9 +28,25 @@
  },
  "entity": {
    "sensor": {
-      "hopfreepowerstart": { "name": "Hour of free power start" },
-      "hopfreepowerend": { "name": "Hour of free power end" }
+      "hop_free_power_start": {
+        "name": "Hour of free power start"
+      },
+      "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": { "hopselector": { "name": "Hour of free power" } }
+    "select": { "hop_selector": { "name": "Hour of free power" } }
  }
 }
--- a/tests/components/electric_kiwi/conftest.py
+++ b/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
--- a/tests/components/electric_kiwi/fixtures/account_balance.json
+++ b/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
+}
--- a/tests/components/electric_kiwi/test_sensor.py
+++ b/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()