update hacs component

custom_components/electric_kiwi/__init__.py

@@ -1,30 +1,30 @@
 """The Electric Kiwi integration."""
+
 from __future__ import annotations
 
 import aiohttp
 from electrickiwi_api import ElectricKiwiApi
-from electrickiwi_api.exceptions import AuthException, ApiException
+from electrickiwi_api.exceptions import ApiException
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 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 (
     ElectricKiwiAccountDataCoordinator,
+    ElectricKiwiConfigEntry,
     ElectricKiwiHOPDataCoordinator,
+    ElectricKiwiRuntimeData,
 )
 
-PLATFORMS: list[Platform] = [
-    Platform.SENSOR,
-    Platform.SELECT,
-]
+PLATFORMS: list[Platform] = [Platform.SELECT, Platform.SENSOR]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ElectricKiwiConfigEntry
+) -> bool:
     """Set up Electric Kiwi from a config entry."""
     implementation = (
         await config_entry_oauth2_flow.async_get_config_entry_implementation(
@@ -44,33 +44,63 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise ConfigEntryNotReady from err
 
     ek_api = ElectricKiwiApi(
-        api.AsyncConfigEntryAuth(aiohttp_client.async_get_clientsession(hass), session)
+        api.ConfigEntryElectricKiwiAuth(
+            aiohttp_client.async_get_clientsession(hass), session
+        )
     )
-    account_coordinator = ElectricKiwiAccountDataCoordinator(hass, ek_api)
-    hop_coordinator = ElectricKiwiHOPDataCoordinator(hass, ek_api)
-
-
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
-        "account_coordinator": account_coordinator,
-        "hop_coordinator": hop_coordinator,
-    }
+    hop_coordinator = ElectricKiwiHOPDataCoordinator(hass, entry, ek_api)
+    account_coordinator = ElectricKiwiAccountDataCoordinator(hass, entry, ek_api)
 
-    # we need to set the client number and connection id
     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
 
+    entry.runtime_data = ElectricKiwiRuntimeData(
+        hop=hop_coordinator, account=account_coordinator
+    )
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(
+    hass: HomeAssistant, entry: ElectricKiwiConfigEntry
+) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+
+
+async def async_migrate_entry(
+    hass: HomeAssistant, config_entry: ElectricKiwiConfigEntry
+) -> bool:
+    """Migrate old entry."""
+    if config_entry.version == 1 and config_entry.minor_version == 1:
+        implementation = (
+            await config_entry_oauth2_flow.async_get_config_entry_implementation(
+                hass, config_entry
+            )
+        )
 
-    return unload_ok
+        session = config_entry_oauth2_flow.OAuth2Session(
+            hass, config_entry, implementation
+        )
+
+        ek_api = ElectricKiwiApi(
+            api.ConfigEntryElectricKiwiAuth(
+                aiohttp_client.async_get_clientsession(hass), session
+            )
+        )
+        try:
+            ek_session = await ek_api.get_active_session()
+        except ApiException:
+            return False
+        unique_id = str(ek_session.data.customer_number)
+        hass.config_entries.async_update_entry(
+            config_entry, unique_id=unique_id, minor_version=2
+        )
+
+    return True

custom_components/electric_kiwi/api.py

@@ -2,17 +2,16 @@
 
 from __future__ import annotations
 
-from typing import cast
-
 from aiohttp import ClientSession
 from electrickiwi_api import AbstractAuth
 
-from homeassistant.helpers import config_entry_oauth2_flow
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow
 
 from .const import API_BASE_URL
 
 
-class AsyncConfigEntryAuth(AbstractAuth):
+class ConfigEntryElectricKiwiAuth(AbstractAuth):
     """Provide Electric Kiwi authentication tied to an OAuth2 based config entry."""
 
     def __init__(
@@ -21,12 +20,29 @@ def __init__(
         oauth_session: config_entry_oauth2_flow.OAuth2Session,
     ) -> None:
         """Initialize Electric Kiwi auth."""
+        # add host when ready for production "https://api.electrickiwi.co.nz" defaults to dev
         super().__init__(websession, API_BASE_URL)
         self._oauth_session = oauth_session
 
     async def async_get_access_token(self) -> str:
         """Return a valid access token."""
-        if not self._oauth_session.valid_token:
-            await self._oauth_session.async_ensure_token_valid()
+        await self._oauth_session.async_ensure_token_valid()
+
+        return str(self._oauth_session.token["access_token"])
+
 
-        return cast(str, self._oauth_session.token["access_token"])
+class ConfigFlowElectricKiwiAuth(AbstractAuth):
+    """Provide Electric Kiwi authentication tied to an OAuth2 based config flow."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        token: str,
+    ) -> None:
+        """Initialize ConfigFlowFitbitApi."""
+        super().__init__(aiohttp_client.async_get_clientsession(hass), API_BASE_URL)
+        self._token = token
+
+    async def async_get_access_token(self) -> str:
+        """Return the token for the Electric Kiwi API."""
+        return self._token

custom_components/electric_kiwi/config_flow.py

@@ -1,14 +1,19 @@
 """Config flow for Electric Kiwi."""
+
 from __future__ import annotations
 
 from collections.abc import Mapping
 import logging
 from typing import Any
 
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.data_entry_flow import FlowResult
+from electrickiwi_api import ElectricKiwiApi
+from electrickiwi_api.exceptions import ApiException
+
+from homeassistant.config_entries import SOURCE_REAUTH, ConfigFlowResult
+from homeassistant.const import CONF_NAME
 from homeassistant.helpers import config_entry_oauth2_flow
 
+from . import api
 from .const import DOMAIN, SCOPE_VALUES
 
 
@@ -17,13 +22,10 @@ class ElectricKiwiOauth2FlowHandler(
 ):
     """Config flow to handle Electric Kiwi OAuth2 authentication."""
 
+    VERSION = 1
+    MINOR_VERSION = 2
     DOMAIN = DOMAIN
 
-    def __init__(self) -> None:
-        """Set up instance."""
-        super().__init__()
-        self._reauth_entry: ConfigEntry | None = None
-
     @property
     def logger(self) -> logging.Logger:
         """Return logger."""
@@ -34,26 +36,41 @@ def extra_authorize_data(self) -> dict[str, Any]:
         """Extra data that needs to be appended to the authorize url."""
         return {"scope": SCOPE_VALUES}
 
-    async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
+    async def async_step_reauth(
+        self, entry_data: Mapping[str, Any]
+    ) -> ConfigFlowResult:
         """Perform reauth upon an API authentication error."""
-        self._reauth_entry = self.hass.config_entries.async_get_entry(
-            self.context["entry_id"]
-        )
         return await self.async_step_reauth_confirm()
 
     async def async_step_reauth_confirm(
         self, user_input: dict[str, Any] | None = None
-    ) -> FlowResult:
+    ) -> ConfigFlowResult:
         """Dialog that informs the user that reauth is required."""
         if user_input is None:
-            return self.async_show_form(step_id="reauth_confirm")
+            return self.async_show_form(
+                step_id="reauth_confirm",
+                description_placeholders={CONF_NAME: self._get_reauth_entry().title},
+            )
         return await self.async_step_user()
 
-    async def async_oauth_create_entry(self, data: dict) -> FlowResult:
+    async def async_oauth_create_entry(self, data: dict) -> ConfigFlowResult:
         """Create an entry for Electric Kiwi."""
-        existing_entry = await self.async_set_unique_id(DOMAIN)
-        if existing_entry:
-            self.hass.config_entries.async_update_entry(existing_entry, data=data)
-            await self.hass.config_entries.async_reload(existing_entry.entry_id)
-            return self.async_abort(reason="reauth_successful")
-        return await super().async_oauth_create_entry(data)
+        ek_api = ElectricKiwiApi(
+            api.ConfigFlowElectricKiwiAuth(self.hass, data["token"]["access_token"])
+        )
+
+        try:
+            session = await ek_api.get_active_session()
+        except ApiException:
+            return self.async_abort(reason="connection_error")
+
+        unique_id = str(session.data.customer_number)
+        await self.async_set_unique_id(unique_id)
+        if self.source == SOURCE_REAUTH:
+            self._abort_if_unique_id_mismatch(reason="wrong_account")
+            return self.async_update_reload_and_abort(
+                self._get_reauth_entry(), data=data
+            )
+
+        self._abort_if_unique_id_configured()
+        return self.async_create_entry(title=unique_id, data=data)

custom_components/electric_kiwi/const.py

@@ -8,14 +8,4 @@
 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"
-
-
-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_SELECT = "hop_select"
-ATTR_EK_HOP_START = "hop_sensor_start"
-ATTR_EK_HOP_END = "hop_sensor_end"
+SCOPE_VALUES = "read_customer_details read_connection_detail read_connection read_billing_address get_bill_address read_billing_frequency read_billing_details read_billing_bills read_billing_bill read_billing_bill_id read_billing_bill_file read_account_running_balance read_customer_account_summary read_consumption_summary download_consumption_file read_consumption_averages get_consumption_averages read_hop_intervals_config read_hop_intervals read_hop_connection read_hop_specific_connection save_hop_connection save_hop_specific_connection read_outage_contact get_outage_contact_info_for_icp read_session read_session_data_login"